#!/usr/bin/env python3 """ agent.py — OSVauco OPAX agent (Jason Vauger). opax_mcp_client er inlina direkte her for Vertex AI Agent Engine-kompatibilitet. """ import asyncio import os import logging import uuid import httpx import json import google.auth import google.auth.transport.requests from typing import Literal from google.adk.agents import Agent from google.adk.tools import FunctionTool from google.adk.runners import Runner from gemini_guard import guarded_run, CircuitOpen, RateLimitExceeded from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.genai import types logger = logging.getLogger(__name__) PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5") LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.5-flash") SUBAGENT_MODEL = os.environ.get("SUBAGENT_MODEL", "gemini-2.5-flash") REASONING_MODEL = os.environ.get("REASONING_MODEL", "gemini-2.5-flash") 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_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" try: from token_logger import log_token_usage, create_bq_table_if_not_exists except ImportError: def log_token_usage(*args, **kwargs): pass def create_bq_table_if_not_exists(): pass try: import sys as _sys, pathlib as _pathlib _sys.path.insert(0, str(_pathlib.Path(__file__).parent.parent.parent / "ml")) from token_budget import trim_context, route_model, TokenBudgetExceeded, budget_summary logger.info(f"[agent] token_budget lastet: {budget_summary()}") except ImportError as _e: logger.warning(f"[agent] token_budget ikke tilgjengelig: {_e}") def trim_context(history, system_prompt="", max_tokens=32000): return history[-10:] if len(history) > 10 else history def route_model(message, mode, flash_model, pro_model): return flash_model if mode != "heavy" else pro_model class TokenBudgetExceeded(Exception): pass # --------------------------------------------------------------------------- # OPAX MCP Client (inlina) # --------------------------------------------------------------------------- _OPAX_MCP_URL = os.environ.get("MCP_SERVER_URL", "https://opax-mcp-zjbqp3prqq-uc.a.run.app") _MCP_SECRET = os.environ.get("MCP_SECRET", "") def _identity_token() -> str: metadata_url = ( "http://metadata.google.internal/computeMetadata/v1/instance" f"/service-accounts/default/identity?audience={_OPAX_MCP_URL}&format=full" ) try: resp = httpx.get(metadata_url, headers={"Metadata-Flavor": "Google"}, timeout=5) if resp.status_code == 200 and resp.text.strip(): return resp.text.strip() except Exception: pass credentials, _ = google.auth.default() credentials.refresh(google.auth.transport.requests.Request()) return credentials.token def _mcp_headers() -> dict: return { "X-MCP-Key": _MCP_SECRET.strip(), "Content-Type": "application/json", } import json ... def _call_tool(tool: str, params: dict = None) -> dict: req_id = str(uuid.uuid4()) resp = httpx.post( f"{_OPAX_MCP_URL}/", headers=_mcp_headers(), json={ "jsonrpc": "2.0", "id": req_id, "method": "tools/call", "params": {"name": tool, "arguments": params or {}}, }, timeout=30, ) resp.raise_for_status() rpc_resp = resp.json() if "error" in rpc_resp: raise Exception(f"Tool call error: {rpc_resp['error']}") return json.loads(rpc_resp["result"]["content"][0]["text"]) def get_billing_summary() -> dict: """Hent billing-oversikt for OSVauco (GCP-kostnader, token-forbruk).""" return _call_tool("get_billing_summary") def get_billing_credits() -> dict: """Hent gjenværende GCP-kreditter og burn-rate.""" return _call_tool("get_billing_credits") def get_billing_anomalies() -> dict: """Sjekk for kostnadanomalier og uventede spiker.""" return _call_tool("get_billing_anomalies") def get_billing_history() -> dict: """Hent historisk kostnadsdata (90 dagar).""" return _call_tool("get_billing_history") def get_billing_budget() -> dict: """Hent gjeldande budsjett for OSVauco.""" return _call_tool("get_billing_budget") def set_billing_budget(amount: float) -> dict: """Sett nytt budsjettgrense. amount er beløpet i USD.""" return _call_tool("set_billing_budget", {"amount": amount}) def create_invite(email: str, company: str = "", tier: str = "starter") -> dict: """Opprett invite-link for ny kunde. tier: starter | guard | shield | enterprise.""" return _call_tool("create_invite", {"email": email, "company": company, "tier": tier}) def list_customers() -> dict: """List alle onboardede kunder og systemtilstand.""" return _call_tool("list_customers") def send_webhook(message: str, url: str = "", title: str = "OPAX varsel") -> dict: """Send webhook-varsling til Slack/Teams/Discord.""" return _call_tool("send_webhook", {"message": message, "url": url, "title": title}) def send_email(to: str, subject: str, body: str = "") -> dict: """Send e-post via SendGrid.""" return _call_tool("send_email", {"to": to, "subject": subject, "body": body}) def get_notify_channels() -> dict: """List konfigurerte varslingskanalar.""" return _call_tool("get_notify_channels") def get_health() -> dict: """Sjekk helsestatus for OPAX-plattformen.""" return _call_tool("get_health") def get_build_status() -> dict: """Hent status på siste Cloud Build-kjøring.""" return _call_tool("get_build_status") def get_state() -> dict: """Hent gjeldande systemtilstand og agent-statistikk for OSVauco.""" return _call_tool("get_state") def get_telemetry() -> dict: """Hent token-forbruk og ytingsdata per modul.""" return _call_tool("get_telemetry") def run_terminal(command: str) -> dict: """Kjør whitelist-kommando i OPAX terminal. Gyldige: health, billing, build, logs, help.""" return _call_tool("run_terminal", {"command": command}) def list_commits(limit: int = 10) -> dict: """List siste commits i OSVauco-repoet på Gitea.""" return _call_tool("list_commits", {"limit": limit}) def get_file(path: str) -> dict: """Hent innhald i ein fil frå Gitea-repoet.""" return _call_tool("get_file", {"path": path}) def list_open_issues() -> dict: """List alle opne issues i OSVauco Gitea-repoet.""" return _call_tool("list_open_issues", {}) def create_issue(title: str, body: str = "") -> dict: """Opprett nytt issue i Gitea-repoet.""" return _call_tool("create_issue", {"title": title, "body": body}) def push_file(path: str, content: str, message: str = "", sha: str = "") -> dict: """Push/oppdater ein fil i Gitea.""" return _call_tool("push_file", {"path": path, "content": content, "message": message, "sha": sha}) OPAX_TOOLS = [ FunctionTool(func=get_billing_summary), FunctionTool(func=get_billing_credits), FunctionTool(func=get_billing_anomalies), FunctionTool(func=get_billing_history), FunctionTool(func=get_billing_budget), FunctionTool(func=set_billing_budget), FunctionTool(func=create_invite), FunctionTool(func=list_customers), FunctionTool(func=send_webhook), FunctionTool(func=send_email), FunctionTool(func=get_notify_channels), FunctionTool(func=get_health), FunctionTool(func=get_build_status), FunctionTool(func=get_state), FunctionTool(func=get_telemetry), FunctionTool(func=run_terminal), FunctionTool(func=list_commits), FunctionTool(func=get_file), FunctionTool(func=list_open_issues), FunctionTool(func=create_issue), FunctionTool(func=push_file), ] # --------------------------------------------------------------------------- # Agent logic # --------------------------------------------------------------------------- def _normalize_mode(mode: str) -> str: 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'.") return mapping[mode] def authorize_mode(user_id: str, mode: str) -> None: 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.") 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} return {"orchestrator": ORCHESTRATOR_MODEL, "subagent": SUBAGENT_MODEL, "reasoning": REASONING_MODEL, "budget_usd": BUDGET_LIGHT} _INSTRUCTION_TEMPLATE = ( "Du er Jason Vauger — OPAX-agenten for Vauco AS. Modus: {mode}. Modell: {model}. " "Du har tilgang til alle OPAX-tools via opax-mcp. " "Tools: list_commits, get_file, push_file, create_issue, list_open_issues (Gitea), " "get_billing_summary, get_billing_credits, get_billing_anomalies, " "get_build_status, get_state, get_telemetry, run_terminal, " "create_invite, list_customers, send_webhook, send_email, get_notify_channels. " "Når brukaren ber om data — KALL ALLTID rett tool før du svarar. " "Svar på norsk (bokmål) med mindre annet er bedt om. " "HITL: ikkje kjør terraform/deploy utan godkjenning frå Chris." ) def build_agent(mode: str = "light", message: str = "") -> tuple: mode = _normalize_mode(mode) models = get_models_for_mode(mode) actual_model = route_model(message=message, mode=mode, flash_model=ORCHESTRATOR_MODEL, pro_model=models["orchestrator"]) instruction = _INSTRUCTION_TEMPLATE.format(mode=mode, model=actual_model) agent = Agent( model=actual_model, name="jason_vauger", description=f"Jason Vauger — OPAX enterprise agent [{mode}]", instruction=instruction, tools=OPAX_TOOLS, ) return agent, actual_model, instruction async def _run_async(message: str, user_id: str, session_id: str, mode: str, caller_type: str = "agent") -> str: mode = _normalize_mode(mode) agent, actual_model, instruction = build_agent(mode=mode, message=message) module_name = f"jason/{mode}" try: trim_context([types.Content(role="user", parts=[types.Part(text=message)])], system_prompt=instruction) except TokenBudgetExceeded as e: return f"⚠️ Token-budsjett overskredet: {e}" 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 = "" input_tokens = 0 output_tokens = 0 request_id = str(uuid.uuid4()) async for event in guarded_run(runner, 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 "" if hasattr(event, "usage_metadata") and event.usage_metadata: um = event.usage_metadata input_tokens += getattr(um, "prompt_token_count", 0) or 0 output_tokens += getattr(um, "candidates_token_count", 0) or 0 if input_tokens > 0 or output_tokens > 0: log_token_usage( agent_name=module_name, model_name=actual_model, input_tokens=input_tokens, output_tokens=output_tokens, request_id=request_id, module_name=module_name, caller_type=caller_type, session_id=session_id, ) logger.info(f"[agent] {mode}/{actual_model} in={input_tokens} out={output_tokens} caller={caller_type} session={session_id}") return final_text def run(message: str, user_id: str = "opax", session_id: str = "default", mode: str = "light", caller_type: str = "agent") -> str: """ Synkron inngang — kun for CLI (if __name__ == '__main__'). Fra FastAPI/uvicorn skal _run_async kalles direkte med await. """ 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, caller_type=caller_type )) if __name__ == "__main__": import sys logging.basicConfig(level=logging.WARNING) query = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX?" mode = sys.argv[2] if len(sys.argv) > 2 else "light" caller_type = sys.argv[3] if len(sys.argv) > 3 else "cli" print(f"Mode: {mode} | Query: {query}") print("-" * 60) print(run(message=query, user_id="opax", mode=mode, caller_type=caller_type))