268 lines
9.6 KiB
Python
268 lines
9.6 KiB
Python
#!/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.
|
|
CG1+CG2: billing endpoints + IAP-beskyttelse.
|
|
CG3: static-mappe serveres fra /static/*.
|
|
AUTH: /auth/login + /auth/callback for klient OAuth2 onboarding.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import pathlib
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "agents", "core-logic"))
|
|
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
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
|
|
from ml.billing_agent import BillingAgent
|
|
from ml.anomaly_detector import AnomalyDetector
|
|
from auth.oauth_flow import get_authorization_url, exchange_code_for_token
|
|
from auth.token_store import save_token
|
|
|
|
app = FastAPI(title="OSVauco OPAX Agent")
|
|
AGENT_ID = "opax-core"
|
|
|
|
# ── IAP-MIDDLEWARE ────────────────────────────────────────────────────────────
|
|
IAP_ENABLED = os.environ.get("IAP_ENABLED", "").lower() == "true"
|
|
IAP_AUDIENCE = os.environ.get("IAP_AUDIENCE", "")
|
|
|
|
|
|
def _verify_iap_jwt(token: str, audience: str) -> dict:
|
|
try:
|
|
from google.auth.transport import requests as google_requests
|
|
from google.oauth2 import id_token
|
|
request = google_requests.Request()
|
|
return id_token.verify_token(token, request, audience=audience,
|
|
certs_url="https://www.gstatic.com/iap/verify/public_key")
|
|
except Exception as exc:
|
|
raise ValueError(f"IAP JWT ugyldig: {exc}") from exc
|
|
|
|
|
|
@app.middleware("http")
|
|
async def iap_guard(request: Request, call_next):
|
|
if IAP_ENABLED and request.url.path.startswith("/billing"):
|
|
token = request.headers.get("X-Goog-IAP-JWT-Assertion", "")
|
|
if not token:
|
|
return JSONResponse(status_code=401,
|
|
content={"error": "IAP-token mangler (X-Goog-IAP-JWT-Assertion)"})
|
|
try:
|
|
_verify_iap_jwt(token, IAP_AUDIENCE)
|
|
except ValueError as exc:
|
|
return JSONResponse(status_code=403, content={"error": str(exc)})
|
|
return await call_next(request)
|
|
|
|
|
|
# ── STATIC FILES ──────────────────────────────────────────────────────────────
|
|
_static_dir = pathlib.Path(__file__).parent / "static"
|
|
if _static_dir.is_dir():
|
|
app.mount("/static", StaticFiles(directory=str(_static_dir), html=True), name="static")
|
|
|
|
|
|
# ── MODELLER ──────────────────────────────────────────────────────────────────
|
|
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")
|
|
|
|
|
|
# ── HEALTH ────────────────────────────────────────────────────────────────────
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
# ── AUTH — OAuth2 klient-onboarding ──────────────────────────────────────────
|
|
@app.get("/auth/login")
|
|
def auth_login(client_id: str):
|
|
"""
|
|
Start OAuth2-flow for en klient.
|
|
Redirect klienten til: GET /auth/login?client_id=<klient-id>
|
|
"""
|
|
try:
|
|
auth_url = get_authorization_url(client_id)
|
|
return RedirectResponse(url=auth_url)
|
|
except Exception as exc:
|
|
return JSONResponse(status_code=500, content={"error": str(exc)})
|
|
|
|
|
|
@app.get("/auth/callback")
|
|
def auth_callback(code: str, state: str):
|
|
"""
|
|
Google redirecter hit etter klient-godkjenning.
|
|
Bytter code mot token og lagrer i Secret Manager.
|
|
"""
|
|
try:
|
|
client_id, token_dict = exchange_code_for_token(code, state)
|
|
save_token(client_id, token_dict)
|
|
return RedirectResponse(url=f"/static/billing-dashboard.html?client_id={client_id}")
|
|
except ValueError as exc:
|
|
return JSONResponse(status_code=400, content={"error": str(exc)})
|
|
except Exception as exc:
|
|
return JSONResponse(status_code=500, content={"error": str(exc)})
|
|
|
|
|
|
# ── BILLING ENDPOINTS (CG1 + CG2) ────────────────────────────────────────────
|
|
@app.get("/billing/summary")
|
|
def billing_summary(client_id: str = ""):
|
|
try:
|
|
return BillingAgent(client_id=client_id or None).get_summary()
|
|
except Exception as exc:
|
|
return JSONResponse(status_code=500, content={"error": str(exc)})
|
|
|
|
|
|
@app.get("/billing/forecast")
|
|
def billing_forecast(client_id: str = ""):
|
|
try:
|
|
return BillingAgent(client_id=client_id or None).get_forecast()
|
|
except Exception as exc:
|
|
return JSONResponse(status_code=500, content={"error": str(exc)})
|
|
|
|
|
|
@app.get("/billing/anomalies")
|
|
def billing_anomalies(client_id: str = ""):
|
|
try:
|
|
return AnomalyDetector(client_id=client_id or None).detect_anomalies()
|
|
except Exception as exc:
|
|
return JSONResponse(status_code=500, content={"error": str(exc)})
|
|
|
|
|
|
# ── AGENT ENDPOINTS ───────────────────────────────────────────────────────────
|
|
@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)
|