OSVauco/agents/core-logic/agent.py

#!/usr/bin/env python3
"""
agent.py — OSVauco OPAX agent (Jason Vauger).

Modes:
  light  — gemini-2.5-flash, $1/task
  heavy  — gemini-2.5-pro,   $3/task

Legacy 'A' / 'A+' normaliseres automatisk.
Autoriserte brukere for heavy: opax, admin

NOTE: opax_mcp_client.py ligg i same mappe (core-logic/) for Vertex AI-kompatibilitet.
"""

import asyncio
import os
import logging
import uuid
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__)

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
    _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

create_bq_table_if_not_exists()


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.")


# —— OPAX-MCP FunctionTools ——
# opax_mcp_client.py er i same mappe — Vertex AI finn den alltid
opax_tools = []
try:
    from opax_mcp_client import get_all_function_tools
    opax_tools = get_all_function_tools()
    logger.info(f"OPAX-MCP tools lastet: {len(opax_tools)} tools")
except Exception as e:
    logger.warning(f"OPAX-MCP tools ikkje tilgjengeleg: {e}")


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 du kan bruke: 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


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)
    models = get_models_for_mode(mode)
    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,
        )

    logger.info(
        f"[agent] {mode}/{actual_model} in={input_tokens} out={output_tokens} "
        f"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:
    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))