#!/usr/bin/env python3 """ agent.py — OSVauco root agent using ADK 1.x with light / heavy mode routing. Modes: light (standard) — gemini-2.0-flash, $1/task hard stop heavy (audit+) — gemini-2.5-pro orchestrator + subagents, gemini-3.5-flash reasoning, $3/task hard stop, multi-agent pipeline enabled Authorized users for heavy: opax, admin Requires: google-adk >= 1.0.0,<2.0.0 google-cloud-aiplatform >= 1.112.0 """ import asyncio import os import logging from typing import Literal from google.adk.agents import Agent from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.genai import types 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 = os.environ.get("RAG_CORPUS", "") # Light mode 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") # Heavy mode 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-3.5-flash") # Budget hard stops (USD per task) BUDGET_LIGHT = float(os.environ.get("BUDGET_A_USD_PER_TASK", "1.0")) BUDGET_HEAVY = float(os.environ.get("HEAVY_MODE_BUDGET_USD_PER_DAY", "3.0")) HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"] Mode = Literal["light", "heavy"] APP_NAME = "opax" def _normalize_mode(mode: str) -> str: """Normalize legacy mode names to light/heavy.""" mapping = {"A": "light", "A+": "heavy", "light": "light", "heavy": "heavy"} if mode not in mapping: raise ValueError(f"Invalid mode '{mode}'. Must be 'light' or 'heavy' (also accepts legacy 'A' / 'A+').") return mapping[mode] # --------------------------------------------------------------------------- # RAG tool # --------------------------------------------------------------------------- rag_tool = None if RAG_CORPUS: try: from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval from vertexai.preview import rag rag_tool = VertexAiRagRetrieval( name="retrieve_knowledge", description="Retrieve relevant documentation and context from the OSVauco knowledge base.", rag_resources=[rag.RagResource(rag_corpus=RAG_CORPUS)], similarity_top_k=10, vector_distance_threshold=0.6, ) logger.info(f"RAG tool initialised with corpus: {RAG_CORPUS}") except ImportError as e: logger.warning(f"VertexAiRagRetrieval not available — running without RAG: {e}") else: logger.warning("RAG_CORPUS env var not set — running without RAG retrieval") # --------------------------------------------------------------------------- # Mode helpers # --------------------------------------------------------------------------- def get_models_for_mode(mode: Mode) -> dict: if mode == "heavy": return { "orchestrator": HEAVY_ORCHESTRATOR, "subagent": HEAVY_SUBAGENT, "reasoning": HEAVY_REASONING, "budget_usd": BUDGET_HEAVY, "multi_agent": True, } return { "orchestrator": ORCHESTRATOR_MODEL, "subagent": SUBAGENT_MODEL, "reasoning": REASONING_MODEL, "budget_usd": BUDGET_LIGHT, "multi_agent": False, } def authorize_mode(user_id: str, mode: str) -> None: """Raise PermissionError if user is not authorized for heavy mode.""" mode = _normalize_mode(mode) if mode == "heavy" and user_id not in HEAVY_MODE_ALLOWED_USERS: raise PermissionError( f"User '{user_id}' is not authorized for heavy mode. " f"Authorized: {HEAVY_MODE_ALLOWED_USERS}" ) # --------------------------------------------------------------------------- # Agent factory # --------------------------------------------------------------------------- def build_agent(mode: str = "light") -> Agent: mode = _normalize_mode(mode) models = get_models_for_mode(mode) mode_label = "heavy" if mode == "heavy" else "light" 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 (light mode) — used by ADK runner and Cloud Run # --------------------------------------------------------------------------- root_agent = build_agent(mode="light") # --------------------------------------------------------------------------- # run() — async core, sync wrapper for Cloud Run /run endpoint # --------------------------------------------------------------------------- async def _run_async( message: str, user_id: str, session_id: str, mode: str, ) -> str: mode = _normalize_mode(mode) agent = build_agent(mode=mode) session_service = InMemorySessionService() session = await session_service.create_session( app_name=APP_NAME, user_id=user_id, session_id=session_id, ) runner = Runner( agent=agent, app_name=APP_NAME, session_service=session_service, ) new_message = types.Content( role="user", parts=[types.Part(text=message)], ) final_text = "" async for event in runner.run_async( user_id=user_id, session_id=session.id, new_message=new_message, ): if event.is_final_response() and event.content and event.content.parts: final_text = event.content.parts[0].text or "" return final_text def run( message: str, user_id: str = "opax", session_id: str = "default", mode: str = "light", ) -> str: """ Handle a single request. Args: message: User message / query. user_id: Caller identity. heavy requires 'opax' or 'admin'. session_id: Session identifier for conversation continuity. mode: 'light' (standard) or 'heavy' (audit+). Legacy values 'A' and 'A+' are accepted and normalized. Returns: Agent response as a string. Raises: PermissionError: If user_id is not authorized for heavy mode. ValueError: If mode is not recognized. """ mode = _normalize_mode(mode) authorize_mode(user_id, mode) return asyncio.run(_run_async( message=message, user_id=user_id, session_id=session_id, mode=mode, )) # --------------------------------------------------------------------------- # CLI smoke test # --------------------------------------------------------------------------- if __name__ == "__main__": import sys logging.basicConfig(level=logging.WARNING) query = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX heavy mode?" requested_mode = sys.argv[2] if len(sys.argv) > 2 else "light" print(f"Mode : {requested_mode}") print(f"Query: {query}") print("-" * 60) print(run(message=query, user_id="opax", mode=requested_mode))