fix: gemini-2.0-flash-001 for light mode, fang ValueError i /run authorize_mode
This commit is contained in:
parent
216fcc5bd4
commit
f711f7be74
|
|
@ -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
|
||||
heavy (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-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
|
||||
|
|
@ -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
|
||||
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")
|
||||
# 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 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")
|
||||
# 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 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]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RAG tool
|
||||
# ---------------------------------------------------------------------------
|
||||
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.",
|
||||
|
|
@ -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")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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:
|
||||
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,
|
||||
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)],
|
||||
)
|
||||
|
||||
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,
|
||||
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
|
||||
|
||||
|
||||
|
|
@ -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,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
mode=mode,
|
||||
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 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)
|
||||
|
|
|
|||
33
main.py
33
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().
|
||||
|
||||
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
|
||||
Modes: light (default, gemini-2.0-flash-001) | heavy (gemini-2.5-pro-001)
|
||||
ML-1: telemetri, state store, DAG-endepunkter.
|
||||
"""
|
||||
|
||||
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):
|
||||
authorize_mode(req.user_id, req.mode)
|
||||
# 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)}
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user