243 lines
12 KiB
Python
243 lines
12 KiB
Python
#!/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 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 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", "global")
|
|
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
|
|
|
|
from token_budget import trim_context, route_model, TokenBudgetExceeded, budget_summary
|
|
except ImportError as _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
|
|
|
|
create_bq_table_if_not_exists()
|
|
|
|
_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:
|
|
try:
|
|
resp = httpx.get(
|
|
f"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?audience={_OPAX_MCP_URL}&format=full",
|
|
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 {"Authorization": f"Bearer {_identity_token()}", "X-MCP-Secret": _MCP_SECRET, "Content-Type": "application/json"}
|
|
|
|
def _call_tool(tool: str, params: dict = None) -> dict:
|
|
resp = httpx.post(f"{_OPAX_MCP_URL}/tools/call", headers=_mcp_headers(), json={"tool": tool, "params": params or {}}, timeout=30)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
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),
|
|
]
|
|
|
|
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. "
|
|
"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,
|
|
vertexai=True,
|
|
project="propane-will-491900-m5",
|
|
location="us-central1",
|
|
)
|
|
return agent, actual_model, instruction
|
|
|
|
root_agent, _, _ = build_agent(mode="light")
|
|
|
|
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 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 ""
|
|
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)
|
|
return final_text
|
|
|
|
def run(message: str, user_id: str = "opax", session_id: str = "default", mode: str = "light", caller_type: str = "agent") -> str:
|
|
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))
|