fix: gemini-2.0-flash-001 for light mode, fang ValueError i /run authorize_mode

agents/core-logic/agent.py

@@ -1,13 +1,12 @@
 #!/usr/bin/env python3
 """
-agent.py — OSVauco root agent using ADK 1.x with light / heavy mode routing.
+agent.py — Cloud Run entrypoint for OSVauco OPAX agent.
 
 Modes:
-  light  (standard) — gemini-2.0-flash, $1/task hard stop
+  light  (standard) — gemini-2.0-flash-001, $1/task hard stop
-  heavy  (audit+)   — gemini-2.5-pro orchestrator + subagents,
+  heavy  (audit+)   — gemini-2.5-pro-001, $3/task hard stop
-                       gemini-3.5-flash reasoning, $3/task hard stop,
-                       multi-agent pipeline enabled
 
+Legacy values 'A' and 'A+' are accepted and normalized.
 Authorized users for heavy: opax, admin
 
 Requires: google-adk >= 1.0.0,<2.0.0
@@ -26,25 +25,20 @@ 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", "")
 
-# Light mode models
+# Light mode — pinned to stable versioned alias
-ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.0-flash")
+ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.0-flash-001")
-SUBAGENT_MODEL     = os.environ.get("SUBAGENT_MODEL",     "gemini-2.0-flash")
+SUBAGENT_MODEL     = os.environ.get("SUBAGENT_MODEL",     "gemini-2.0-flash-001")
-REASONING_MODEL    = os.environ.get("REASONING_MODEL",    "gemini-2.0-flash")
+REASONING_MODEL    = os.environ.get("REASONING_MODEL",    "gemini-2.0-flash-001")
 
-# Heavy mode models
+# Heavy mode
-HEAVY_ORCHESTRATOR = os.environ.get("HEAVY_ORCHESTRATOR_MODEL", "gemini-2.5-pro")
+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")
+HEAVY_SUBAGENT     = os.environ.get("HEAVY_SUBAGENT_MODEL",     "gemini-2.5-pro-001")
-HEAVY_REASONING    = os.environ.get("HEAVY_REASONING_MODEL",    "gemini-3.5-flash")
+HEAVY_REASONING    = os.environ.get("HEAVY_REASONING_MODEL",    "gemini-2.5-flash-001")
 
-# Budget hard stops (USD per task)
 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"))
 
@@ -55,23 +49,30 @@ 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+').")
+        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:
-# RAG tool
+    """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.",
@@ -86,10 +87,6 @@ else:
     logger.warning("RAG_CORPUS env var not set — running without RAG retrieval")
 
 
-# ---------------------------------------------------------------------------
-# Mode helpers
-# ---------------------------------------------------------------------------
-
 def get_models_for_mode(mode: Mode) -> dict:
     if mode == "heavy":
         return {
@@ -108,25 +105,10 @@ def get_models_for_mode(mode: Mode) -> dict:
     }
 
 
-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 heavy mode. "
-            f"Authorized: {HEAVY_MODE_ALLOWED_USERS}"
-        )
-
-
-# ---------------------------------------------------------------------------
-# Agent factory
-# ---------------------------------------------------------------------------
-
 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']}. "
@@ -134,7 +116,6 @@ def build_agent(mode: str = "light") -> Agent:
         "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",
@@ -144,53 +125,24 @@ def build_agent(mode: str = "light") -> Agent:
     )
 
 
-# ---------------------------------------------------------------------------
-# Default root agent (light mode) — used by ADK runner and Cloud Run
-# ---------------------------------------------------------------------------
-
 root_agent = build_agent(mode="light")
 
 
-# ---------------------------------------------------------------------------
+async def _run_async(message: str, user_id: str, session_id: str, mode: str) -> str:
-# run() — async core, sync wrapper for Cloud Run /run endpoint
-# ---------------------------------------------------------------------------
-
-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,
+        app_name=APP_NAME, user_id=user_id, session_id=session_id,
-        user_id=user_id,
-        session_id=session_id,
     )
-
+    runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service)
-    runner = Runner(
+    new_message = types.Content(role="user", parts=[types.Part(text=message)])
-        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,
+        user_id=user_id, session_id=session.id, new_message=new_message,
-        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
 
 
@@ -200,45 +152,18 @@ def run(
     session_id: str = "default",
     mode: str = "light",
 ) -> str:
-    """
-    Handle a single request.
-
-    Args:
-        message:    User message / query.
-        user_id:    Caller identity. heavy requires 'opax' or 'admin'.
-        session_id: Session identifier for conversation continuity.
-        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 heavy mode.
-        ValueError:      If mode is not recognized.
-    """
     mode = _normalize_mode(mode)
     authorize_mode(user_id, mode)
-
     return asyncio.run(_run_async(
-        message=message,
+        message=message, user_id=user_id, session_id=session_id, mode=mode,
-        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 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)

main.py

@@ -1,15 +1,8 @@
 #!/usr/bin/env python3
 """
 main.py — Cloud Run entrypoint for OSVauco OPAX agent.
-Exposes a FastAPI HTTP API that wraps agent.run().
+Modes: light (default, gemini-2.0-flash-001) | heavy (gemini-2.5-pro-001)
-
+ML-1: telemetri, state store, DAG-endepunkter.
-ML-1 changes:
-  - /run now records timing + outcome via ml.telemetry
-  - /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
@@ -42,7 +35,7 @@ class DagRequest(BaseModel):
     user_id: str = "opax"
     session_id: str = "default"
     mode: str = "light"
-    scheduler: str = Field("threads", description="Dask scheduler: synchronous | threads | processes")
+    scheduler: str = Field("threads")
 
 
 @app.get("/health")
@@ -52,7 +45,13 @@ def health():
 
 @app.post("/run")
 def run_agent(req: RunRequest):
+    # Validate mode and authorization before any work — returns 400/403 on bad input
+    try:
         authorize_mode(req.user_id, req.mode)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except PermissionError as e:
+        raise HTTPException(status_code=403, detail=str(e))
 
     store = get_store()
     start = time.monotonic()
@@ -66,18 +65,13 @@ def run_agent(req: RunRequest):
             session_id=req.session_id,
             mode=req.mode,
         )
-    except PermissionError as e:
-        raise HTTPException(status_code=403, detail=str(e))
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         error_msg = str(e)
         raise HTTPException(status_code=500, detail=error_msg)
     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"
+        model = "gemini-2.0-flash-001" if req.mode in ("light", "A") else "gemini-2.5-pro-001"
-
         log_agent_call(
             agent_id=AGENT_ID,
             input_payload={"message": req.message, "mode": req.mode},
@@ -101,6 +95,8 @@ def run_agent(req: RunRequest):
 def run_dag(req: DagRequest):
     try:
         authorize_mode(req.user_id, req.mode)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
     except PermissionError as e:
         raise HTTPException(status_code=403, detail=str(e))
 
@@ -150,17 +146,14 @@ def run_dag(req: DagRequest):
 def get_state():
     return get_store().snapshot()
 
-
 @app.get("/state/agents")
 def list_agents():
     return {"agents": get_store().list_agents()}
 
-
 @app.get("/state/aggregate/{key}")
 def aggregate_key(key: str):
     return get_store().aggregate(key)
 
-
 @app.get("/telemetry/history")
 def telemetry_history(limit: int = 50):
     return {"history": get_store().history(limit=limit)}