revert: gjenopprett original main.py fra fungerende commit 6981359
This commit is contained in:
parent
e85ba6de4f
commit
d63ff5bab5
664
main.py
664
main.py
|
|
@ -1,7 +1,663 @@
|
|||
import sys
|
||||
#!/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
|
||||
|
||||
# /app er WORKDIR, agents/core-logic ligg der etter COPY . .
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'agents', 'core-logic'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "agents", "core-logic"))
|
||||
|
||||
from app import app # noqa: F401
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List
|
||||
import datetime
|
||||
from cachetools import cached, TTLCache
|
||||
|
||||
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 agents.aws_billing_agent import AWSBillingAgent
|
||||
from ml.anomaly_detector import AnomalyDetector
|
||||
from auth.token_store import save_token
|
||||
|
||||
|
||||
from google.cloud import firestore
|
||||
import firebase_admin
|
||||
from firebase_admin import credentials, messaging
|
||||
import base64
|
||||
import json
|
||||
from functools import wraps
|
||||
from starlette.middleware.sessions import SessionMiddleware
|
||||
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
|
||||
from authlib.integrations.starlette_client import OAuth
|
||||
import sendgrid
|
||||
from sendgrid.helpers.mail import Mail
|
||||
|
||||
AGENT_ID = "osvauco-opax"
|
||||
app = FastAPI(
|
||||
title="OSVauco OPAX Agent",
|
||||
description="Agent for OSVauco-OPAX platform.",
|
||||
version="0.1.0",
|
||||
)
|
||||
|
||||
# Paths exempt from IAP enforcement (health/readiness probes reach Cloud Run
|
||||
# directly without the IAP-injected x-goog-authenticated-user-email header).
|
||||
IAP_EXEMPT_PATHS = {"/health", "/healthz", "/readiness", "/liveness"}
|
||||
|
||||
@app.middleware("http")
|
||||
async def require_iap(request: Request, call_next):
|
||||
if request.url.path in IAP_EXEMPT_PATHS:
|
||||
return await call_next(request)
|
||||
if not request.headers.get("x-goog-authenticated-user-email"):
|
||||
return Response(status_code=401, content="Unauthorized")
|
||||
return await call_next(request)
|
||||
|
||||
# ── FIREBASE & FIRESTORE INIT ───────────────────────────────────────────────
|
||||
db = None
|
||||
try:
|
||||
# When running in Google Cloud, ADC will be used automatically.
|
||||
# For local dev, `gcloud auth application-default login` is required.
|
||||
firebase_admin.initialize_app()
|
||||
db = firestore.Client()
|
||||
print("Firestore client initialized successfully.")
|
||||
except Exception as e:
|
||||
# The app can still run in a limited mode without Firestore.
|
||||
# Endpoints that depend on `db` will return a 500 error.
|
||||
print(f"WARNING: Firestore client failed to initialize: {e}", file=sys.stderr)
|
||||
|
||||
# ── AUTH & SESSION ────────────────────────────────────────────────────────────
|
||||
# Add session middleware for storing auth state
|
||||
app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="*")
|
||||
app.add_middleware(SessionMiddleware, secret_key=os.environ.get("SESSION_SECRET"))
|
||||
|
||||
# Configure Authlib's OAuth client
|
||||
oauth = OAuth()
|
||||
oauth.register(
|
||||
name='google',
|
||||
client_id=os.environ.get("GOOGLE_CLIENT_ID"),
|
||||
client_secret=os.environ.get("GOOGLE_CLIENT_SECRET"),
|
||||
server_metadata_url='https://accounts.google.com/.well-known/openid-configuration',
|
||||
client_kwargs={
|
||||
'scope': 'openid email profile'
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@app.get('/auth/login')
|
||||
async def login(request: Request):
|
||||
"""Redirects user to Google's OAuth 2.0 login page."""
|
||||
redirect_uri = "https://opax.vauco.no/auth/callback"
|
||||
return await oauth.google.authorize_redirect(request, redirect_uri)
|
||||
|
||||
@app.get('/auth/callback', name='auth')
|
||||
async def auth(request: Request):
|
||||
"""
|
||||
Handles the callback from Google's OAuth.
|
||||
Stores user info in session and redirects to the billing dashboard.
|
||||
"""
|
||||
token = await oauth.google.authorize_access_token(request)
|
||||
user = token.get('userinfo')
|
||||
if user:
|
||||
request.session['user'] = dict(user)
|
||||
return RedirectResponse(url='/static/billing-dashboard.html')
|
||||
|
||||
|
||||
@app.get("/auth/me")
|
||||
async def me(request: Request):
|
||||
"""Returns the authenticated user's information."""
|
||||
user = request.session.get("user")
|
||||
if not user:
|
||||
raise HTTPException(status_code=401, detail="Not authenticated")
|
||||
return JSONResponse(user)
|
||||
|
||||
|
||||
@app.get('/auth/logout')
|
||||
async def logout(request: Request):
|
||||
"""Clears the user session and logs them out."""
|
||||
request.session.pop('user', None)
|
||||
return RedirectResponse(url='/static/billing-dashboard.html')
|
||||
|
||||
|
||||
# Allowed emails for login
|
||||
ALLOWED_EMAILS = [email.strip() for email in os.environ.get("ALLOWED_EMAILS", "").split(",") if email.strip()]
|
||||
ALERT_EMAIL = os.environ.get("ALERT_EMAIL")
|
||||
|
||||
|
||||
def require_auth(func):
|
||||
"""Decorator to protect endpoints that require authentication."""
|
||||
@wraps(func)
|
||||
async def wrapper(request: Request, *args, **kwargs):
|
||||
user = request.session.get('user')
|
||||
if not user:
|
||||
return JSONResponse(status_code=401, content={"error": "Not authenticated"})
|
||||
|
||||
if ALLOWED_EMAILS and user.get('email') not in ALLOWED_EMAILS:
|
||||
return JSONResponse(status_code=403, content={"error": "Email not allowed"})
|
||||
|
||||
return await func(request, *args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
|
||||
# ── 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")
|
||||
|
||||
class PushSubscription(BaseModel):
|
||||
token: str
|
||||
budget_nok: float
|
||||
|
||||
class BudgetWebhookPayload(BaseModel):
|
||||
message: dict
|
||||
subscription: str
|
||||
|
||||
class BudgetUpdateRequest(BaseModel):
|
||||
budget: float
|
||||
|
||||
# ── ROOT LANDING PAGE ───────────────────────────────────────────────────
|
||||
@app.get("/")
|
||||
def root():
|
||||
"""Serves the main landing page."""
|
||||
return FileResponse("static/opax.html")
|
||||
|
||||
|
||||
# ── HEALTH ────────────────────────────────────────────────────────────────────
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
@app.get("/manifest.json", include_in_schema=False)
|
||||
def manifest():
|
||||
return FileResponse("static/manifest.json")
|
||||
|
||||
|
||||
@app.get("/sw.js", include_in_schema=False)
|
||||
def service_worker():
|
||||
return FileResponse("static/sw.js")
|
||||
|
||||
|
||||
# ── AUTH — User Authentication ──────────────────────────────────────────
|
||||
@app.get('/admin')
|
||||
@require_auth
|
||||
async def admin_panel(request: Request):
|
||||
return FileResponse("static/admin.html")
|
||||
|
||||
@app.post('/admin/create-customer')
|
||||
@require_auth
|
||||
async def create_customer(request: Request):
|
||||
import subprocess, shlex
|
||||
data = await request.json()
|
||||
customer_name = data.get("customer_name", "").strip()
|
||||
project_id = data.get("project_id", "").strip()
|
||||
billing_account_id = data.get("billing_account_id", "").strip()
|
||||
alert_email = data.get("alert_email", "").strip()
|
||||
container_image = data.get("container_image", "").strip()
|
||||
region = data.get("region", "europe-north1").strip()
|
||||
|
||||
if not all([customer_name, project_id, billing_account_id, alert_email, container_image]):
|
||||
raise HTTPException(status_code=400, detail="Alle felt er påkrevd")
|
||||
|
||||
import re
|
||||
if not re.match(r'^[a-z0-9_-]+$', customer_name):
|
||||
raise HTTPException(status_code=400, detail="customer_name kan kun inneholde a-z, 0-9, - og _")
|
||||
|
||||
customer_dir = f"infrastructure/terraform/customers/{customer_name}"
|
||||
template_dir = "infrastructure/terraform/customers/_template"
|
||||
|
||||
import os, shutil
|
||||
if os.path.exists(customer_dir):
|
||||
raise HTTPException(status_code=409, detail=f"Kunde {customer_name} eksisterer allerede")
|
||||
|
||||
shutil.copytree(template_dir, customer_dir)
|
||||
|
||||
tfvars_content = f'''customer_id = "{customer_name}"
|
||||
project_id = "{project_id}"
|
||||
region = "{region}"
|
||||
billing_account_id = "{billing_account_id}"
|
||||
alert_email = "{alert_email}"
|
||||
billing_viewer_emails = ["chris.christiansen@vauco.no", "jason.vauger@vauco.no"]
|
||||
container_image = "{container_image}"
|
||||
'''
|
||||
with open(f"{customer_dir}/terraform.tfvars", "w") as f:
|
||||
f.write(tfvars_content)
|
||||
|
||||
try:
|
||||
init = subprocess.run(
|
||||
["terraform", f"-chdir={customer_dir}", "init", "-no-color"],
|
||||
capture_output=True, text=True, timeout=120
|
||||
)
|
||||
if init.returncode != 0:
|
||||
shutil.rmtree(customer_dir)
|
||||
raise HTTPException(status_code=500, detail=f"terraform init feilet: {init.stderr[-500:]}")
|
||||
|
||||
apply = subprocess.run(
|
||||
["terraform", f"-chdir={customer_dir}", "apply", "-auto-approve", "-no-color"],
|
||||
capture_output=True, text=True, timeout=600
|
||||
)
|
||||
if apply.returncode != 0:
|
||||
raise HTTPException(status_code=500, detail=f"terraform apply feilet: {apply.stderr[-500:]}")
|
||||
|
||||
return {"status": "ok", "customer": customer_name, "project_id": project_id}
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
raise HTTPException(status_code=504, detail="Terraform tok for lang tid (>10 min)")
|
||||
|
||||
|
||||
# ── BILLING ENDPOINTS (CG1 + CG2) ────────────────────────────────────────────
|
||||
|
||||
async def get_budget(request: Request):
|
||||
if not db:
|
||||
return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
|
||||
try:
|
||||
doc_ref = db.collection("settings").document("budget")
|
||||
doc = doc_ref.get()
|
||||
if doc.exists:
|
||||
return JSONResponse(content={"budget": doc.to_dict().get("limit", 500)})
|
||||
else:
|
||||
return JSONResponse(content={"budget": 500}) # Default
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||
|
||||
async def set_budget(request: Request, payload: BudgetUpdateRequest):
|
||||
if not db:
|
||||
return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
|
||||
try:
|
||||
doc_ref = db.collection("settings").document("budget")
|
||||
doc_ref.set({"limit": payload.budget})
|
||||
return JSONResponse(content={"status": "ok", "budget": payload.budget})
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||
|
||||
app.add_api_route("/billing/budget", endpoint=require_auth(get_budget), methods=["GET"])
|
||||
app.add_api_route("/billing/budget", endpoint=require_auth(set_budget), methods=["POST"])
|
||||
|
||||
|
||||
@app.post("/billing/email-report")
|
||||
async def trigger_email_report():
|
||||
"""
|
||||
Generates and sends a daily cost summary email.
|
||||
"""
|
||||
if not ALERT_EMAIL or not os.environ.get("SENDGRID_API_KEY"):
|
||||
err_msg = "Email reporting is not configured. Missing ALERT_EMAIL or SENDGRID_API_KEY."
|
||||
print(err_msg, file=sys.stderr)
|
||||
# Return success to scheduler to prevent retries for config errors
|
||||
return JSONResponse(status_code=200, content={"status": "not_configured", "error": err_msg})
|
||||
|
||||
try:
|
||||
# Fetch data
|
||||
billing_agent = BillingAgent()
|
||||
summary_data_response = billing_agent.get_summary()
|
||||
summary_data = summary_data_response.get("summary", [])
|
||||
forecast_data = billing_agent.get_forecast()
|
||||
|
||||
# Get budget
|
||||
budget_response = await get_budget(None)
|
||||
budget = json.loads(budget_response.body).get("budget", 500)
|
||||
|
||||
# Format email content
|
||||
mtd_cost = forecast_data.get("month_to_date_cost", 0)
|
||||
daily_avg = forecast_data.get("daily_average_last_7_days", 0)
|
||||
top_services = summary_data[:3]
|
||||
|
||||
service_list_html = "".join([f"<li>{s['service']}: kr {s.get('total_cost', 0):.2f}</li>" for s in top_services])
|
||||
|
||||
html_content = f"""
|
||||
<h3>CostGuard Daglig Oppsummering - {datetime.date.today()}</h3>
|
||||
<p><strong>Kostnad MTD:</strong> kr {mtd_cost:.2f}</p>
|
||||
<p><strong>Daglig snitt (7d):</strong> kr {daily_avg:.2f}</p>
|
||||
<p><strong>Budsjettstatus:</strong> kr {mtd_cost:.2f} av kr {budget:.2f}</p>
|
||||
<p><strong>Topp 3 GCP-tjenester:</strong></p>
|
||||
<ul>{service_list_html}</ul>
|
||||
<p>--<br>CostGuard by Vauco</p>
|
||||
"""
|
||||
|
||||
message = Mail(
|
||||
from_email='costguard@osvauco.no',
|
||||
to_emails=ALERT_EMAIL,
|
||||
subject=f'CostGuard Daglig Oppsummering {datetime.date.today()}',
|
||||
html_content=html_content
|
||||
)
|
||||
|
||||
sg = sendgrid.SendGridAPIClient(os.environ.get('SENDGRID_API_KEY'))
|
||||
response = sg.send(message)
|
||||
|
||||
print(f"Email report sent successfully, message ID: {response.headers.get('X-Message-Id')}")
|
||||
return JSONResponse(content={"status": "ok", "message_id": response.headers.get('X-Message-Id')})
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error sending email report: {e}", file=sys.stderr)
|
||||
# Return 500 to indicate a transient failure that scheduler might retry
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||
|
||||
|
||||
@app.post("/billing/snapshot")
|
||||
async def create_daily_snapshot(request: Request):
|
||||
"""
|
||||
Called by Cloud Scheduler. Fetches current billing data and saves a snapshot to Firestore.
|
||||
This endpoint should be secured by Cloud Scheduler's authentication (e.g., OIDC token).
|
||||
"""
|
||||
# ... (snapshot logic remains the same) ...
|
||||
|
||||
# After successfully creating the snapshot, trigger the email report.
|
||||
email_response = await trigger_email_report()
|
||||
|
||||
# Log the outcome of the email sending but don't let it fail the snapshot creation.
|
||||
if email_response.status_code != 200:
|
||||
print(f"Daily snapshot created, but failed to send email report. Status: {email_response.status_code}", file=sys.stderr)
|
||||
|
||||
return JSONResponse(content={"status": "ok", "snapshot_id": date_str})
|
||||
|
||||
|
||||
async def get_history(request: Request):
|
||||
if not db:
|
||||
return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
|
||||
try:
|
||||
end_date = datetime.date.today()
|
||||
start_date = end_date - datetime.timedelta(days=90)
|
||||
|
||||
docs = db.collection("daily_snapshots") \
|
||||
.where("created_at", ">=", start_date.isoformat()) \
|
||||
.order_by("created_at", direction=firestore.Query.DESCENDING) \
|
||||
.limit(90) \
|
||||
.stream()
|
||||
|
||||
history = [{"date": doc.id, **doc.to_dict()} for doc in docs]
|
||||
|
||||
# The query is descending, so we reverse to get chronological order
|
||||
history.reverse()
|
||||
|
||||
return JSONResponse(content={"history": history})
|
||||
except Exception as e:
|
||||
print(f"Error fetching history: {e}", file=sys.stderr)
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||
|
||||
app.add_api_route("/billing/history", endpoint=require_auth(get_history), methods=["GET"])
|
||||
|
||||
|
||||
# Note: The @require_auth decorator is a function, not an async function,
|
||||
# so we can't use it directly on FastAPI routes like this.
|
||||
# Instead, we will wrap the functions manually for now.
|
||||
# A more robust solution would use Depends.
|
||||
|
||||
async def authenticated_billing_summary(request: Request):
|
||||
try:
|
||||
return BillingAgent().get_summary()
|
||||
except Exception as exc:
|
||||
return JSONResponse(status_code=500, content={"error": str(exc)})
|
||||
app.add_api_route("/billing/summary", endpoint=require_auth(authenticated_billing_summary), methods=["GET"])
|
||||
|
||||
|
||||
|
||||
async def authenticated_billing_forecast(request: Request):
|
||||
try:
|
||||
return BillingAgent().get_forecast()
|
||||
except Exception as exc:
|
||||
return JSONResponse(status_code=500, content={"error": str(exc)})
|
||||
app.add_api_route("/billing/forecast", endpoint=require_auth(authenticated_billing_forecast), methods=["GET"])
|
||||
|
||||
|
||||
async def authenticated_billing_anomalies(request: Request):
|
||||
try:
|
||||
return AnomalyDetector().detect_anomalies()
|
||||
except Exception as exc:
|
||||
return JSONResponse(status_code=500, content={"error": str(exc)})
|
||||
app.add_api_route("/billing/anomalies", endpoint=require_auth(authenticated_billing_anomalies), methods=["GET"])
|
||||
|
||||
|
||||
@app.get("/billing-dashboard")
|
||||
def billing_dashboard_view(request: Request):
|
||||
"""Serves the main dashboard and injects Firebase config."""
|
||||
# This endpoint is now deprecated in favor of the static file,
|
||||
# but we keep it in case there are old links.
|
||||
# The new auth flow redirects to the static file directly.
|
||||
return RedirectResponse(url='/static/billing-dashboard.html')
|
||||
|
||||
@app.post("/billing/subscribe")
|
||||
def subscribe_for_push(sub: PushSubscription):
|
||||
if not db:
|
||||
raise HTTPException(status_code=500, detail="Firestore is not configured")
|
||||
try:
|
||||
doc_ref = db.collection("push_subscribers").document(sub.token)
|
||||
doc_ref.set({"budget_nok": sub.budget_nok, "subscribed_at": firestore.SERVER_TIMESTAMP})
|
||||
return {"status": "ok"}
|
||||
except Exception as e:
|
||||
print(f"Error subscribing token: {e}", file=sys.stderr)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/billing/budget-webhook")
|
||||
async def budget_webhook(payload: BudgetWebhookPayload):
|
||||
if not db:
|
||||
print("Webhook called but Firestore is not configured. Aborting.", file=sys.stderr)
|
||||
# Return 200 to prevent Pub/Sub from retrying
|
||||
return {"status": "error", "detail": "Firestore not configured"}
|
||||
try:
|
||||
data = base64.b64decode(payload.message.get("data", "")).decode("utf-8")
|
||||
data_json = json.loads(data)
|
||||
|
||||
cost_amount = data_json.get("costAmount", 0)
|
||||
budget_amount = data_json.get("budgetAmount", 0)
|
||||
|
||||
print(f"Webhook received: cost={cost_amount}, budget={budget_amount}")
|
||||
|
||||
if budget_amount > 0 and (cost_amount / budget_amount) > 0.8:
|
||||
percent_used = round((cost_amount / budget_amount) * 100)
|
||||
|
||||
subscribers_ref = db.collection("push_subscribers")
|
||||
subscribers = subscribers_ref.stream()
|
||||
|
||||
tokens = [subscriber.id for subscriber in subscribers]
|
||||
|
||||
if not tokens:
|
||||
print("Budget threshold exceeded, but no push subscribers found.")
|
||||
return {"status": "no subscribers"}
|
||||
|
||||
notification = messaging.Notification(
|
||||
title="⚠️ CostGuard Varsel",
|
||||
body=f"Du har brukt {percent_used}% av budsjett (kr {int(cost_amount)} av kr {int(budget_amount)})"
|
||||
)
|
||||
|
||||
message = messaging.MulticastMessage(
|
||||
tokens=tokens,
|
||||
notification=notification
|
||||
)
|
||||
|
||||
response = messaging.send_multicast(message)
|
||||
print(f"Sent push notification to {response.success_count} subscribers.")
|
||||
else:
|
||||
print("Budget threshold not exceeded, no notification sent.")
|
||||
|
||||
return {"status": "processed"}
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in budget webhook: {e}", file=sys.stderr)
|
||||
# Return 200 to prevent Pub/Sub from retrying
|
||||
return {"status": "error", "detail": str(e)}
|
||||
|
||||
|
||||
@app.get("/billing/live")
|
||||
@cached(TTLCache(maxsize=1, ttl=3600))
|
||||
async def billing_live(request: Request):
|
||||
"""
|
||||
Henter live faktureringsdata fra Cloud Billing Budgets API.
|
||||
Dette gir raskere, men mindre detaljert, data enn BigQuery-eksporten.
|
||||
Fallback til BigQuery ved feil.
|
||||
"""
|
||||
# This endpoint should also be protected.
|
||||
user = request.session.get('user')
|
||||
if not user:
|
||||
return JSONResponse(status_code=401, content={"error": "Not authenticated"})
|
||||
if ALLOWED_EMAILS and user.get('email') not in ALLOWED_EMAILS:
|
||||
return JSONResponse(status_code=403, content={"error": "Email not allowed"})
|
||||
|
||||
# Fallback to BigQuery forecast
|
||||
print("Billing Budget API not implemented, falling back to BigQuery", file=sys.stderr)
|
||||
bq_forecast = BillingAgent().get_forecast()
|
||||
bq_forecast["data_source"] = "BigQuery Fallback"
|
||||
return bq_forecast
|
||||
|
||||
|
||||
# ── AWS BILLING ENDPOINTS (CG6) ──────────────────────────────────────────────
|
||||
|
||||
async def authenticated_aws_billing_summary(request: Request):
|
||||
try:
|
||||
return AWSBillingAgent().get_summary()
|
||||
except Exception as exc:
|
||||
return JSONResponse(status_code=500, content={"error": str(exc)})
|
||||
app.add_api_route("/billing/aws/summary", endpoint=require_auth(authenticated_aws_billing_summary), methods=["GET"])
|
||||
|
||||
async def authenticated_aws_billing_forecast(request: Request):
|
||||
try:
|
||||
return AWSBillingAgent().get_forecast()
|
||||
except Exception as exc:
|
||||
return JSONResponse(status_code=500, content={"error": str(exc)})
|
||||
app.add_api_route("/billing/aws/forecast", endpoint=require_auth(authenticated_aws_billing_forecast), methods=["GET"])
|
||||
|
||||
|
||||
# ── 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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user