chrischristiansen-glitch
7686fc7623
- token_logger.py: nytt schema med caller_type (agent|cli|browser|cron|api), module_name (jason/opax, costguard, threadstone, cli/gemini-tui, billing_agent), session_id - agent.py: agent_name → strukturert label (module_name + mode), caller_type satt til 'agent' - ROADMAP.md: CG4 Token Intelligence-fase lagt til (CG4a–CG4d) under NEXT
237 lines
8.2 KiB
Python
237 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
agent.py — OSVauco OPAX agent.
|
|
|
|
Modes:
|
|
light — gemini-2.5-flash, $1/task
|
|
heavy — gemini-2.5-pro, $3/task
|
|
|
|
Legacy 'A' / 'A+' normaliseres automatisk.
|
|
Authoriserte brukere for heavy: opax, admin
|
|
|
|
Gjeldende modell-tilgjengelighet (mai 2026):
|
|
gemini-2.0-flash-001 kun for eksisterende kunder — bruk IKKE.
|
|
gemini-2.5-flash / gemini-2.5-pro krever GOOGLE_CLOUD_LOCATION=global.
|
|
"""
|
|
|
|
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")
|
|
RAG_CORPUS = os.environ.get("RAG_CORPUS", "")
|
|
|
|
# Light mode — gemini-2.5-flash (GA, 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 mode — gemini-2.5-pro (GA, global)
|
|
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"
|
|
|
|
# Token logger — feiler stille, stopper aldri agent
|
|
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
|
|
|
|
# Opprett BQ-tabell ved oppstart (idempotent, feiler stille)
|
|
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' "
|
|
f"(also accepts legacy 'A' / 'A+')."
|
|
)
|
|
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. "
|
|
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: {RAG_CORPUS}")
|
|
except ImportError as e:
|
|
logger.warning(f"VertexAiRagRetrieval ikke tilgjengelig: {e}")
|
|
else:
|
|
logger.warning("RAG_CORPUS ikke satt — kjører uten RAG")
|
|
|
|
|
|
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)
|
|
instruction = (
|
|
f"Du er OPAX — OSVauco AI-agent (kallenavn: Jason). Modus: {mode}. "
|
|
f"Du kjører på Vertex AI i region us-central1. "
|
|
f"Aktiv modell: {models['orchestrator']} "
|
|
f"(heavy = gemini-2.5-pro, light = gemini-2.5-flash). "
|
|
f"Oppgi ALDRI andre modellnavn enn disse. "
|
|
f"Budsjettgrense per oppgave: ${models['budget_usd']}. "
|
|
"Bruk retrieve_knowledge for å hente dokumentasjon. "
|
|
"Sanntids GCP-kostnader hentes fra /billing/summary, /billing/forecast "
|
|
"og /billing/anomalies. Hvis disse mangler data, er det fordi BigQuery "
|
|
"billing_export (OQ-15) ikke er aktivert ennå — si dette ærlig, ikke gjett tall. "
|
|
"HITL-grenser: du kan ikke kjøre terraform apply / onboarde kunder, "
|
|
"selge Vauco OS eller OPAX, eller passere godkjenningsporter uten Chris. "
|
|
"Svar på norsk (bokmål) med mindre annet er bedt om."
|
|
)
|
|
return Agent(
|
|
model=models["orchestrator"],
|
|
name="opax_agent",
|
|
description=f"OPAX — OSVauco enterprise agent [{mode}]",
|
|
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,
|
|
caller_type: str = "agent",
|
|
) -> str:
|
|
mode = _normalize_mode(mode)
|
|
agent = build_agent(mode=mode)
|
|
models = get_models_for_mode(mode)
|
|
module_name = f"jason/{mode}" # CG4: 'jason/light' eller 'jason/heavy'
|
|
|
|
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
|
|
|
|
# Logg med full CG4-labeling — feiler stille
|
|
if input_tokens > 0 or output_tokens > 0:
|
|
log_token_usage(
|
|
agent_name=module_name, # bakoverkompatibelt
|
|
model_name=models["orchestrator"],
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
request_id=request_id,
|
|
module_name=module_name, # CG4: 'jason/light' | 'jason/heavy'
|
|
caller_type=caller_type, # CG4: agent|browser|cli|cron|api
|
|
session_id=session_id, # CG4: for sesjon-gruppering
|
|
)
|
|
|
|
return final_text
|
|
|
|
|
|
def run(
|
|
message: str,
|
|
user_id: str = "opax",
|
|
session_id: str = "default",
|
|
mode: str = "light",
|
|
caller_type: str = "agent",
|
|
) -> str:
|
|
"""
|
|
Kjør OPAX-agenten.
|
|
|
|
caller_type-verdier:
|
|
'agent' — programmatisk kall fra server/app.py
|
|
'browser' — direkte fra browser-chat (opax.vauco.no)
|
|
'cli' — lokal kjøring fra terminalen
|
|
'cron' — bakgrunnsjobb
|
|
'api' — ekstern REST-kall
|
|
"""
|
|
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} | Caller: {caller_type} | Query: {query}")
|
|
print("-" * 60)
|
|
print(run(message=query, user_id="opax", mode=mode, caller_type=caller_type))
|