feat(ml-1): wire ML DAG + state store + telemetry into /run endpoint

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