239 lines
7.7 KiB
Python
239 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
agent.py — OSVauco root agent using ADK 1.x with A / A+ mode routing.
|
|
|
|
Modes:
|
|
A (standard) — gemini-2.0-flash-001, $1/task hard stop
|
|
A+ (audit+) — gemini-2.5-pro orchestrator + subagents,
|
|
gemini-3.5-flash reasoning, $3/task hard stop,
|
|
multi-agent pipeline enabled
|
|
|
|
Authorized users for A+: 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", "")
|
|
|
|
# Standard (A) models
|
|
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")
|
|
|
|
# A+ (audit+) 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") # NOT 2.5-flash
|
|
|
|
# Budget hard stops (USD per task)
|
|
BUDGET_A = float(os.environ.get("BUDGET_A_USD_PER_TASK", "1.0"))
|
|
BUDGET_APLUS = float(os.environ.get("HEAVY_MODE_BUDGET_USD_PER_DAY", "3.0"))
|
|
|
|
HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]
|
|
|
|
Mode = Literal["A", "A+"]
|
|
APP_NAME = "opax"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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:
|
|
"""Return model config dict for the given mode."""
|
|
if mode == "A+":
|
|
return {
|
|
"orchestrator": HEAVY_ORCHESTRATOR,
|
|
"subagent": HEAVY_SUBAGENT,
|
|
"reasoning": HEAVY_REASONING,
|
|
"budget_usd": BUDGET_APLUS,
|
|
"multi_agent": True,
|
|
}
|
|
return {
|
|
"orchestrator": ORCHESTRATOR_MODEL,
|
|
"subagent": SUBAGENT_MODEL,
|
|
"reasoning": REASONING_MODEL,
|
|
"budget_usd": BUDGET_A,
|
|
"multi_agent": False,
|
|
}
|
|
|
|
|
|
def authorize_mode(user_id: str, mode: Mode) -> None:
|
|
"""Raise PermissionError if user is not authorized for A+ mode."""
|
|
if mode == "A+" and user_id not in HEAVY_MODE_ALLOWED_USERS:
|
|
raise PermissionError(
|
|
f"User '{user_id}' is not authorized for A+ mode. "
|
|
f"Authorized: {HEAVY_MODE_ALLOWED_USERS}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Agent factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_agent(mode: Mode = "A") -> Agent:
|
|
"""Build and return an ADK Agent configured for the given mode."""
|
|
models = get_models_for_mode(mode)
|
|
mode_label = "A+ (audit+)" if mode == "A+" else "A (standard)"
|
|
|
|
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 (A mode) — used by ADK runner and Cloud Run
|
|
# ---------------------------------------------------------------------------
|
|
|
|
root_agent = build_agent(mode="A")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# run() — async core, sync wrapper for Cloud Run /run endpoint
|
|
# ---------------------------------------------------------------------------
|
|
|
|
async def _run_async(
|
|
message: str,
|
|
user_id: str,
|
|
session_id: str,
|
|
mode: Mode,
|
|
) -> str:
|
|
"""Async implementation using ADK 1.x Runner pattern."""
|
|
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 = "A",
|
|
) -> str:
|
|
"""
|
|
Handle a single request.
|
|
|
|
Args:
|
|
message: User message / query.
|
|
user_id: Caller identity. A+ requires 'opax' or 'admin'.
|
|
session_id: Session identifier for conversation continuity.
|
|
mode: 'A' (standard) or 'A+' (audit+).
|
|
|
|
Returns:
|
|
Agent response as a string.
|
|
|
|
Raises:
|
|
PermissionError: If user_id is not authorized for A+ mode.
|
|
ValueError: If mode is not 'A' or 'A+'.
|
|
"""
|
|
if mode not in ("A", "A+"):
|
|
raise ValueError(f"Invalid mode '{mode}'. Must be 'A' or 'A+'")
|
|
|
|
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 A+ mode?"
|
|
requested_mode = sys.argv[2] if len(sys.argv) > 2 else "A"
|
|
|
|
print(f"Mode : {requested_mode}")
|
|
print(f"Query: {query}")
|
|
print("-" * 60)
|
|
print(run(message=query, user_id="opax", mode=requested_mode))
|