diff --git a/agents/core-logic/agent.py b/agents/core-logic/agent.py index 0796e35..8f84031 100644 --- a/agents/core-logic/agent.py +++ b/agents/core-logic/agent.py @@ -1,14 +1,14 @@ #!/usr/bin/env python3 """ -agent.py — OSVauco root agent using ADK 1.x with A / A+ mode routing. +agent.py — OSVauco root agent using ADK 1.x with light / heavy mode routing. Modes: - A (standard) — gemini-2.0-flash, $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 + light (standard) — gemini-2.0-flash, $1/task hard stop + heavy (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 +Authorized users for heavy: opax, admin Requires: google-adk >= 1.0.0,<2.0.0 google-cloud-aiplatform >= 1.112.0 @@ -34,26 +34,34 @@ 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 +# Light mode models ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.0-flash") SUBAGENT_MODEL = os.environ.get("SUBAGENT_MODEL", "gemini-2.0-flash") REASONING_MODEL = os.environ.get("REASONING_MODEL", "gemini-2.0-flash") -# A+ (audit+) models +# Heavy mode 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 +HEAVY_REASONING = os.environ.get("HEAVY_REASONING_MODEL", "gemini-3.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")) +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["A", "A+"] +Mode = Literal["light", "heavy"] APP_NAME = "opax" +def _normalize_mode(mode: str) -> str: + """Normalize legacy mode names to light/heavy.""" + 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' (also accepts legacy 'A' / 'A+').") + return mapping[mode] + + # --------------------------------------------------------------------------- # RAG tool # --------------------------------------------------------------------------- @@ -83,29 +91,29 @@ else: # --------------------------------------------------------------------------- def get_models_for_mode(mode: Mode) -> dict: - """Return model config dict for the given mode.""" - if mode == "A+": + if mode == "heavy": return { "orchestrator": HEAVY_ORCHESTRATOR, "subagent": HEAVY_SUBAGENT, "reasoning": HEAVY_REASONING, - "budget_usd": BUDGET_APLUS, + "budget_usd": BUDGET_HEAVY, "multi_agent": True, } return { "orchestrator": ORCHESTRATOR_MODEL, "subagent": SUBAGENT_MODEL, "reasoning": REASONING_MODEL, - "budget_usd": BUDGET_A, + "budget_usd": BUDGET_LIGHT, "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: +def authorize_mode(user_id: str, mode: str) -> None: + """Raise PermissionError if user is not authorized for heavy mode.""" + 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 A+ mode. " + f"User '{user_id}' is not authorized for heavy mode. " f"Authorized: {HEAVY_MODE_ALLOWED_USERS}" ) @@ -114,10 +122,10 @@ def authorize_mode(user_id: str, mode: Mode) -> None: # Agent factory # --------------------------------------------------------------------------- -def build_agent(mode: Mode = "A") -> Agent: - """Build and return an ADK Agent configured for the given mode.""" +def build_agent(mode: str = "light") -> Agent: + mode = _normalize_mode(mode) models = get_models_for_mode(mode) - mode_label = "A+ (audit+)" if mode == "A+" else "A (standard)" + mode_label = "heavy" if mode == "heavy" else "light" instruction = ( f"Du er OPAX — OSVauco AI-agent. " @@ -137,10 +145,10 @@ def build_agent(mode: Mode = "A") -> Agent: # --------------------------------------------------------------------------- -# Default root agent (A mode) — used by ADK runner and Cloud Run +# Default root agent (light mode) — used by ADK runner and Cloud Run # --------------------------------------------------------------------------- -root_agent = build_agent(mode="A") +root_agent = build_agent(mode="light") # --------------------------------------------------------------------------- @@ -151,9 +159,9 @@ async def _run_async( message: str, user_id: str, session_id: str, - mode: Mode, + mode: str, ) -> str: - """Async implementation using ADK 1.x Runner pattern.""" + mode = _normalize_mode(mode) agent = build_agent(mode=mode) session_service = InMemorySessionService() @@ -190,27 +198,26 @@ def run( message: str, user_id: str = "opax", session_id: str = "default", - mode: str = "A", + mode: str = "light", ) -> str: """ Handle a single request. Args: message: User message / query. - user_id: Caller identity. A+ requires 'opax' or 'admin'. + user_id: Caller identity. heavy requires 'opax' or 'admin'. session_id: Session identifier for conversation continuity. - mode: 'A' (standard) or 'A+' (audit+). + mode: 'light' (standard) or 'heavy' (audit+). + Legacy values 'A' and 'A+' are accepted and normalized. 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+'. + PermissionError: If user_id is not authorized for heavy mode. + ValueError: If mode is not recognized. """ - if mode not in ("A", "A+"): - raise ValueError(f"Invalid mode '{mode}'. Must be 'A' or 'A+'.") - + mode = _normalize_mode(mode) authorize_mode(user_id, mode) return asyncio.run(_run_async( @@ -229,8 +236,8 @@ 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" + 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}") diff --git a/main.py b/main.py index 4af44ad..8bb3681 100644 --- a/main.py +++ b/main.py @@ -8,68 +8,51 @@ ML-1 changes: - /run pushes last_result + last_duration_s to AgentStateStore - /run/dag accepts a list of messages and runs them as a parallel Dask DAG - /state and /telemetry endpoints expose live ML observability + +Modes: light (default) | heavy """ import os import sys import time -# Ensure agents/core-logic is on the path sys.path.insert(0, os.path.join(os.path.dirname(__file__), "agents", "core-logic")) from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field -from typing import List, Optional +from typing import List -from agent import run, authorize_mode # agents/core-logic/agent.py -from ml import ( - build_agent_dag, - execute_dag, - get_store, - log_agent_call, -) +from agent import run, authorize_mode +from ml import build_agent_dag, execute_dag, get_store, log_agent_call from ml.telemetry import log_dag_execution app = FastAPI(title="OSVauco OPAX Agent") - AGENT_ID = "opax-core" -# --------------------------------------------------------------------------- -# Request / Response models -# --------------------------------------------------------------------------- - class RunRequest(BaseModel): message: str user_id: str = "opax" session_id: str = "default" - mode: str = "A" + mode: str = "light" class DagRequest(BaseModel): messages: List[str] = Field(..., description="Liste av meldinger som kjøres parallelt") user_id: str = "opax" session_id: str = "default" - mode: str = "A" + mode: str = "light" scheduler: str = Field("threads", description="Dask scheduler: synchronous | threads | processes") -# --------------------------------------------------------------------------- -# Health -# --------------------------------------------------------------------------- - @app.get("/health") def health(): return {"status": "ok"} -# --------------------------------------------------------------------------- -# /run — enkelt agent-kall med telemetri + state store (ML-1) -# --------------------------------------------------------------------------- - @app.post("/run") def run_agent(req: RunRequest): - authorize_mode(req.user_id, req.mode) # raises PermissionError if unauthorized + authorize_mode(req.user_id, req.mode) store = get_store() start = time.monotonic() @@ -93,20 +76,18 @@ def run_agent(req: RunRequest): finally: duration = round(time.monotonic() - start, 3) success = error_msg is None + model = "gemini-2.0-flash" if req.mode in ("light", "A") else "gemini-2.5-pro" - # ML-1: logg til JSONL-telemetri log_agent_call( agent_id=AGENT_ID, input_payload={"message": req.message, "mode": req.mode}, output=response, - model_used="gemini-2.0-flash" if req.mode == "A" else "gemini-2.5-pro", + model_used=model, mode=req.mode, duration_s=duration, success=success, error=error_msg, ) - - # ML-1: push til Parameter Server state store store.push(AGENT_ID, "last_duration_s", duration) store.push(AGENT_ID, "last_mode", req.mode) store.push(AGENT_ID, "last_success", success) @@ -116,24 +97,14 @@ def run_agent(req: RunRequest): return {"response": response} -# --------------------------------------------------------------------------- -# /run/dag — parallell Dask DAG over flere meldinger (ML-1) -# --------------------------------------------------------------------------- - @app.post("/run/dag") def run_dag(req: DagRequest): - """ - Kjører en liste av meldinger som parallelle Dask-tasks. - Hver melding behandles av en separat agent-wrapper. - Returnerer alle svar med timing og suksess-status. - """ try: authorize_mode(req.user_id, req.mode) except PermissionError as e: raise HTTPException(status_code=403, detail=str(e)) def make_agent_fn(msg: str, idx: int): - """Lukker over msg + idx for å lage en unik callable per melding.""" def _agent_fn(payload: dict): return run( message=msg, @@ -144,12 +115,7 @@ def run_dag(req: DagRequest): _agent_fn.__name__ = f"opax-dag-{idx}" return _agent_fn - payload = { - "user_id": req.user_id, - "session_id": req.session_id, - "mode": req.mode, - } - + payload = {"user_id": req.user_id, "session_id": req.session_id, "mode": req.mode} agent_fns = [make_agent_fn(msg, i) for i, msg in enumerate(req.messages)] agent_ids = [f"opax-dag-{i}" for i in range(len(req.messages))] @@ -158,13 +124,8 @@ def run_dag(req: DagRequest): results = execute_dag(tasks, scheduler=req.scheduler) total_dur = round(time.monotonic() - start, 3) - # ML-1: logg DAG-kjøring og push aggregert state store = get_store() - log_dag_execution( - dag_id=req.session_id, - agent_results=results, - total_duration_s=total_dur, - ) + log_dag_execution(dag_id=req.session_id, agent_results=results, total_duration_s=total_dur) store.push(AGENT_ID, "last_dag_total_duration_s", total_dur) store.push(AGENT_ID, "last_dag_agent_count", len(results)) store.push(AGENT_ID, "last_dag_success_count", sum(1 for r in results if r["success"])) @@ -185,42 +146,26 @@ def run_dag(req: DagRequest): } -# --------------------------------------------------------------------------- -# /state — live Parameter Server state snapshot (ML-1 observability) -# --------------------------------------------------------------------------- - @app.get("/state") def get_state(): - """Returnerer nåværende AgentStateStore snapshot.""" return get_store().snapshot() @app.get("/state/agents") def list_agents(): - """Returnerer alle agent-IDer som har pushet state.""" return {"agents": get_store().list_agents()} @app.get("/state/aggregate/{key}") def aggregate_key(key: str): - """Aggregerer alle agenter sine verdier for en gitt nøkkel.""" return get_store().aggregate(key) -# --------------------------------------------------------------------------- -# /telemetry — historikk-log (ML-1 observability) -# --------------------------------------------------------------------------- - @app.get("/telemetry/history") def telemetry_history(limit: int = 50): - """Returnerer siste N state-operasjoner fra AgentStateStore historikk.""" return {"history": get_store().history(limit=limit)} -# --------------------------------------------------------------------------- -# Startup -# --------------------------------------------------------------------------- - if __name__ == "__main__": import uvicorn port = int(os.environ.get("PORT", 8080))