OSVauco/main.py

#!/usr/bin/env python3
"""
main.py — Cloud Run entrypoint, OSVauco OPAX.
Modes: light (gemini-2.5-flash) | heavy (gemini-2.5-pro)
ML-1: telemetri, state store, DAG.
"""

import os
import sys
import time

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

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"


class RunRequest(BaseModel):
    message: str
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "light"


class DagRequest(BaseModel):
    messages: List[str] = Field(...)
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "light"
    scheduler: str = Field("threads")


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/run")
def run_agent(req: RunRequest):
    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()
    error_msg = None
    response = None

    try:
        response = run(
            message=req.message,
            user_id=req.user_id,
            session_id=req.session_id,
            mode=req.mode,
        )
    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.5-flash" if req.mode in ("light", "A") else "gemini-2.5-pro"
        log_agent_call(
            agent_id=AGENT_ID,
            input_payload={"message": req.message, "mode": req.mode},
            output=response,
            model_used=model,
            mode=req.mode,
            duration_s=duration,
            success=success,
            error=error_msg,
        )
        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}


@app.post("/run/dag")
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))

    def make_agent_fn(msg: str, idx: int):
        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)

    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)
        ],
    }


@app.get("/state")
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)}


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8080))
    uvicorn.run(app, host="0.0.0.0", port=port)