feat: add A/A+ mode routing to agent.py (OPAX)
This commit is contained in:
parent
60c7aff04d
commit
497509370e
|
|
@ -1,27 +1,60 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
agent.py — OSVauco root agent using ADK 1.x.
|
||||
agent.py — OSVauco root agent using ADK 1.x with A / A+ mode routing.
|
||||
|
||||
Modes:
|
||||
A (standard) — gemini-2.0-flash, $1/task hard stop
|
||||
A+ (audit+) — gemini-2.5-pro orchestrator, $3/task hard stop,
|
||||
multi-agent pipeline enabled
|
||||
|
||||
Authorized users for A+: opax, admin
|
||||
|
||||
Requires: google-adk >= 1.0.0,<2.0.0
|
||||
google-cloud-aiplatform >= 1.112.0
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Literal
|
||||
|
||||
from google.adk.agents import Agent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
|
||||
LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
|
||||
# RAG corpus resource name — injected via Secret Manager as env-var RAG_CORPUS
|
||||
RAG_CORPUS = os.environ.get("RAG_CORPUS", "")
|
||||
|
||||
# Standard (A) models
|
||||
ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.0-flash")
|
||||
SUBAGENT_MODEL = os.environ.get("SUBAGENT_MODEL", "gemini-2.0-flash")
|
||||
REASONING_MODEL = os.environ.get("REASONING_MODEL", "gemini-2.0-flash")
|
||||
|
||||
# A+ (audit+) models
|
||||
HEAVY_ORCHESTRATOR = os.environ.get("HEAVY_ORCHESTRATOR_MODEL", "gemini-2.5-pro")
|
||||
HEAVY_SUBAGENT = os.environ.get("HEAVY_SUBAGENT_MODEL", "gemini-2.5-pro")
|
||||
HEAVY_REASONING = os.environ.get("HEAVY_REASONING_MODEL", "gemini-2.5-flash")
|
||||
|
||||
# Budget hard stops (USD per task)
|
||||
BUDGET_A = float(os.environ.get("BUDGET_A_USD_PER_TASK", "1.0"))
|
||||
BUDGET_APLUS = float(os.environ.get("HEAVY_MODE_BUDGET_USD_PER_DAY", "3.0"))
|
||||
|
||||
HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]
|
||||
|
||||
Mode = Literal["A", "A+"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RAG tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
rag_tool = None
|
||||
if RAG_CORPUS:
|
||||
try:
|
||||
# ADK 1.x: VertexAiRagRetrieval lives in google.adk.tools.retrieval
|
||||
from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
|
||||
from vertexai.preview import rag
|
||||
|
||||
|
|
@ -35,20 +68,126 @@ if RAG_CORPUS:
|
|||
logger.info(f"RAG tool initialised with corpus: {RAG_CORPUS}")
|
||||
except ImportError as e:
|
||||
logger.warning(f"VertexAiRagRetrieval not available — running without RAG: {e}")
|
||||
rag_tool = None
|
||||
else:
|
||||
logger.warning("RAG_CORPUS env var not set — running without RAG retrieval")
|
||||
|
||||
# Root agent
|
||||
root_agent = Agent(
|
||||
model="gemini-2.5-flash",
|
||||
name="osvauco_root",
|
||||
description="OSVauco enterprise agent for propane-will-491900-m5",
|
||||
instruction=(
|
||||
"You are OSVauco, a GCP knowledge and workflow agent. "
|
||||
"Use the retrieve_knowledge tool to answer questions from the knowledge base. "
|
||||
"Always prefer grounded, documented answers over speculation. "
|
||||
"Always respond in Norwegian (Bokmål) regardless of the language used in the query."
|
||||
),
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mode helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def get_models_for_mode(mode: Mode) -> dict:
|
||||
"""Return model config dict for the given mode."""
|
||||
if mode == "A+":
|
||||
return {
|
||||
"orchestrator": HEAVY_ORCHESTRATOR,
|
||||
"subagent": HEAVY_SUBAGENT,
|
||||
"reasoning": HEAVY_REASONING,
|
||||
"budget_usd": BUDGET_APLUS,
|
||||
"multi_agent": True,
|
||||
}
|
||||
return {
|
||||
"orchestrator": ORCHESTRATOR_MODEL,
|
||||
"subagent": SUBAGENT_MODEL,
|
||||
"reasoning": REASONING_MODEL,
|
||||
"budget_usd": BUDGET_A,
|
||||
"multi_agent": False,
|
||||
}
|
||||
|
||||
|
||||
def authorize_mode(user_id: str, mode: Mode) -> None:
|
||||
"""Raise PermissionError if user is not authorized for A+ mode."""
|
||||
if mode == "A+" and user_id not in HEAVY_MODE_ALLOWED_USERS:
|
||||
raise PermissionError(
|
||||
f"User '{user_id}' is not authorized for A+ mode. "
|
||||
f"Authorized: {HEAVY_MODE_ALLOWED_USERS}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def build_agent(mode: Mode = "A") -> Agent:
|
||||
"""Build and return an ADK Agent configured for the given mode."""
|
||||
models = get_models_for_mode(mode)
|
||||
mode_label = "A+ (audit+)" if mode == "A+" else "A (standard)"
|
||||
|
||||
instruction = (
|
||||
f"Du er OPAX — OSVauco AI-agent. "
|
||||
f"Modus: {mode_label}. Budsjettgrense per oppgave: ${models['budget_usd']}. "
|
||||
"Bruk retrieve_knowledge-verktøyet for å hente dokumentasjon og kontekst. "
|
||||
"Foretrekk alltid dokumenterte svar fremfor spekulasjon. "
|
||||
"Svar alltid på norsk (bokmål) med mindre brukeren eksplisitt ber om et annet språk."
|
||||
)
|
||||
|
||||
return Agent(
|
||||
model=models["orchestrator"],
|
||||
name="opax_agent",
|
||||
description=f"OPAX — OSVauco enterprise agent [{mode_label}]",
|
||||
instruction=instruction,
|
||||
tools=[rag_tool] if rag_tool else [],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Default root agent (A mode) — used by ADK runner and Cloud Run
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
root_agent = build_agent(mode="A")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# run() — programmatic entry point for Cloud Run /run endpoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run(
|
||||
message: str,
|
||||
user_id: str = "opax",
|
||||
session_id: str = "default",
|
||||
mode: str = "A",
|
||||
) -> str:
|
||||
"""
|
||||
Handle a single request.
|
||||
|
||||
Args:
|
||||
message: User message / query.
|
||||
user_id: Caller identity. A+ requires 'opax' or 'admin'.
|
||||
session_id: Session identifier for conversation continuity.
|
||||
mode: 'A' (standard) or 'A+' (audit+).
|
||||
|
||||
Returns:
|
||||
Agent response as a string.
|
||||
|
||||
Raises:
|
||||
PermissionError: If user_id is not authorized for the requested mode.
|
||||
ValueError: If mode is not 'A' or 'A+'.
|
||||
"""
|
||||
if mode not in ("A", "A+"):
|
||||
raise ValueError(f"Invalid mode '{mode}'. Must be 'A' or 'A+'.")
|
||||
|
||||
authorize_mode(user_id, mode)
|
||||
|
||||
agent = build_agent(mode=mode)
|
||||
response = agent.run(
|
||||
message=message,
|
||||
session_id=session_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
return response.text if hasattr(response, "text") else str(response)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI smoke test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
query = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX A+ mode?"
|
||||
requested_mode = sys.argv[2] if len(sys.argv) > 2 else "A"
|
||||
|
||||
print(f"Mode : {requested_mode}")
|
||||
print(f"Query: {query}")
|
||||
print("-" * 60)
|
||||
print(run(message=query, user_id="opax", mode=requested_mode))
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user