OSVauco/agents/core-logic/agent.py

#!/usr/bin/env python3
"""
agent.py — Cloud Run entrypoint for OSVauco OPAX agent.

Modes:
  light  (standard) — gemini-2.0-flash-001, $1/task hard stop
  heavy  (audit+)   — gemini-2.5-pro-001, $3/task hard stop

Legacy values 'A' and 'A+' are accepted and normalized.
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__)

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 — pinned to stable versioned alias
ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.0-flash-001")
SUBAGENT_MODEL     = os.environ.get("SUBAGENT_MODEL",     "gemini-2.0-flash-001")
REASONING_MODEL    = os.environ.get("REASONING_MODEL",    "gemini-2.0-flash-001")

# Heavy mode
HEAVY_ORCHESTRATOR = os.environ.get("HEAVY_ORCHESTRATOR_MODEL", "gemini-2.5-pro-001")
HEAVY_SUBAGENT     = os.environ.get("HEAVY_SUBAGENT_MODEL",     "gemini-2.5-pro-001")
HEAVY_REASONING    = os.environ.get("HEAVY_REASONING_MODEL",    "gemini-2.5-flash-001")

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:
    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' "
            f"(also accepts legacy 'A' / 'A+')."
        )
    return mapping[mode]


def authorize_mode(user_id: str, mode: str) -> None:
    """Normalize mode first, then check authorization. Raises ValueError or PermissionError."""
    mode = _normalize_mode(mode)  # raises ValueError for unknown modes
    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}"
        )


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


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 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 [],
    )


root_agent = build_agent(mode="light")


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:
    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,
    ))


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