From 497509370e180f733d1936f7db17f0c0956746bc Mon Sep 17 00:00:00 2001 From: chrischristiansen-glitch Date: Sun, 24 May 2026 16:34:34 +0200 Subject: [PATCH] feat: add A/A+ mode routing to agent.py (OPAX) --- agents/core-logic/agent.py | 175 +++++++++++++++++++++++++++++++++---- 1 file changed, 157 insertions(+), 18 deletions(-) diff --git a/agents/core-logic/agent.py b/agents/core-logic/agent.py index 7525dff..b678fe6 100644 --- a/agents/core-logic/agent.py +++ b/agents/core-logic/agent.py @@ -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." - ), - tools=[rag_tool] if rag_tool else [], -) + +# --------------------------------------------------------------------------- +# 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))