diff --git a/main.py b/main.py index 332d25c..4af44ad 100644 --- a/main.py +++ b/main.py @@ -2,20 +2,42 @@ """ 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 """ 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 -from agent import run # agents/core-logic/agent.py +from pydantic import BaseModel, Field +from typing import List, Optional + +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 ml.telemetry import log_dag_execution app = FastAPI(title="OSVauco OPAX Agent") +AGENT_ID = "opax-core" + + +# --------------------------------------------------------------------------- +# Request / Response models +# --------------------------------------------------------------------------- class RunRequest(BaseModel): message: str @@ -24,13 +46,36 @@ class RunRequest(BaseModel): mode: str = "A" +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" + 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 + + store = get_store() + start = time.monotonic() + error_msg = None + response = None + try: response = run( message=req.message, @@ -38,14 +83,143 @@ def run_agent(req: RunRequest): session_id=req.session_id, mode=req.mode, ) - return {"response": response} 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: - raise HTTPException(status_code=500, detail=str(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 + # 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", + 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) + if response: + store.push(AGENT_ID, "last_result_preview", response[:200]) + + 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, + user_id=payload["user_id"], + session_id=f"{payload['session_id']}-dag-{idx}", + mode=payload["mode"], + ) + _agent_fn.__name__ = f"opax-dag-{idx}" + return _agent_fn + + 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))] + + start = time.monotonic() + tasks = build_agent_dag(agent_fns, payload, agent_ids) + 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, + ) + 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"])) + + return { + "total_duration_s": total_dur, + "results": [ + { + "index": i, + "message": req.messages[i], + "response": r["result"], + "success": r["success"], + "duration_s": r["duration_s"], + "error": r["error"], + } + for i, r in enumerate(results) + ], + } + + +# --------------------------------------------------------------------------- +# /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