aetherflow_api.py

from flask import Flask, request, jsonify
import traceback

# Import core credit management helpers
from aetherflow_core import (
    register_agent,
    deduct_credits,
    add_credits,
    get_balance,
    authorize_agent,
    analytics_summary,
    analytics_agents,
    save_policy,
    get_policy,
    list_policies,
    update_policy,
)

app = Flask(__name__)

"""AetherFlow Credits API — pay-as-you-go credit metering for AI agents"""

# ---------------------------------------------------------------------------
# Auth helper — returns (agent_id, None) on success or (None, error_response)
# ---------------------------------------------------------------------------
def _require_auth():
    api_key = request.headers.get('X-API-KEY') or request.headers.get('Authorization', '').replace('Bearer ', '')
    if not api_key:
        return None, (jsonify({"error": "missing_api_key", "hint": "Include X-API-KEY header"}), 401)
    agent_id = authorize_agent(api_key)
    if not agent_id:
        return None, (jsonify({"error": "invalid_api_key"}), 403)
    return agent_id, None


# ---------------------------------------------------------------------------
# Public endpoints
# ---------------------------------------------------------------------------
@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "ok", "service": "aetherflow-credits"})


@app.route('/agents/register', methods=['POST'])
def register():
    """Register a new agent — returns agent_id + api_key. Public endpoint."""
    data = request.get_json(force=True) or {}
    agent_id = data.get('agent_id')
    initial_credits = data.get('initial_credits', 1000)
    if not agent_id:
        return jsonify({"error": "agent_id is required"}), 400
    try:
        result = register_agent(agent_id, initial_credits)
        return jsonify(result)
    except Exception:
        traceback.print_exc()
        return jsonify({"error": "internal_error"}), 500


# ---------------------------------------------------------------------------
# Protected agent endpoints — require X-API-KEY
# ---------------------------------------------------------------------------
@app.route('/agents/deduct', methods=['POST'])
def deduct():
    agent_id, err = _require_auth()
    if err:
        return err
    data = request.get_json(force=True) or {}
    # Allow body agent_id override; fall back to authenticated agent
    target_id = data.get('agent_id', agent_id)
    amount = data.get('amount')
    description = data.get('description', '')
    if amount is None:
        return jsonify({"error": "amount is required"}), 400
    try:
        policy_id = data.get('policy_id')
        msg = deduct_credits(target_id, amount, description, policy_id)
        balance = get_balance(target_id)
        return jsonify({"message": msg, "balance": balance})
    except Exception:
        traceback.print_exc()
        return jsonify({"error": "internal_error"}), 500


@app.route('/agents/add', methods=['POST'])
def add():
    agent_id, err = _require_auth()
    if err:
        return err
    data = request.get_json(force=True) or {}
    target_id = data.get('agent_id', agent_id)
    amount = data.get('amount')
    description = data.get('description', '')
    if amount is None:
        return jsonify({"error": "amount is required"}), 400
    try:
        msg = add_credits(target_id, amount, description)
        balance = get_balance(target_id)
        return jsonify({"message": msg, "balance": balance})
    except Exception:
        traceback.print_exc()
        return jsonify({"error": "internal_error"}), 500


@app.route('/agents/balance', methods=['GET'])
def balance():
    agent_id, err = _require_auth()
    if err:
        return err
    target_id = request.args.get('agent_id', agent_id)
    try:
        bal = get_balance(target_id)
        return jsonify(bal)
    except Exception:
        traceback.print_exc()
        return jsonify({"error": "internal_error"}), 500


# ---------------------------------------------------------------------------
# Policy endpoints — protected
# ---------------------------------------------------------------------------
@app.route('/policies/create', methods=['POST'])
def policies_create():
    _, err = _require_auth()
    if err:
        return err
    policy = request.get_json(force=True) or {}
    result = save_policy(policy)
    return jsonify(result)


@app.route('/policies/<policy_id>', methods=['GET'])
def policies_get(policy_id):
    _, err = _require_auth()
    if err:
        return err
    policy = get_policy(policy_id)
    if not policy:
        return jsonify({"error": "policy_not_found"}), 404
    return jsonify({"policy": policy})


@app.route('/policies', methods=['GET'])
def policies_list():
    _, err = _require_auth()
    if err:
        return err
    return jsonify({"policies": list_policies()})


@app.route('/policies/<policy_id>', methods=['PATCH'])
def policies_patch(policy_id):
    _, err = _require_auth()
    if err:
        return err
    patch = request.get_json(force=True) or {}
    result = update_policy(policy_id, patch)
    return jsonify(result)


@app.route('/policies/<policy_id>/enable-auto-topup', methods=['POST'])
def policies_enable_autotop(policy_id):
    _, err = _require_auth()
    if err:
        return err
    policy = get_policy(policy_id)
    if not policy:
        return jsonify({"error": "policy_not_found"}), 404
    patch = {"auto_topup": {"enabled": True}}
    return jsonify(update_policy(policy_id, patch))


@app.route('/policies/<policy_id>/disable-auto-topup', methods=['POST'])
def policies_disable_autotop(policy_id):
    _, err = _require_auth()
    if err:
        return err
    policy = get_policy(policy_id)
    if not policy:
        return jsonify({"error": "policy_not_found"}), 404
    patch = {"auto_topup": {"enabled": False}}
    return jsonify(update_policy(policy_id, patch))


# ---------------------------------------------------------------------------
# Analytics endpoints — protected
# ---------------------------------------------------------------------------
@app.route('/analytics/summary', methods=['GET'])
def analytics_sum():
    _, err = _require_auth()
    if err:
        return err
    return jsonify(analytics_summary())


@app.route('/analytics/agents', methods=['GET'])
def analytics_agents_route():
    _, err = _require_auth()
    if err:
        return err
    return jsonify(analytics_agents())


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080, debug=False)