id: litellm-proxy-budget-and-model-fallback-gate
namespace: company.aiops

inputs:
  - id: primary_provider
    type: STRING
    defaults: "openai/gpt-4o"
    description: Primary LLM model routed through the enterprise LiteLLM proxy.

  - id: fallback_provider
    type: STRING
    defaults: "anthropic/claude-3-5-sonnet"
    description: Failover model provider invoked when primary latency or error rate
      degrades.

  - id: daily_budget_usd
    type: FLOAT
    defaults: 150.0
    description: Daily expenditure budget threshold in USD before triggering spend alerts.

tasks:
  - id: check_provider_health_and_spend
    type: io.kestra.plugin.scripts.python.Script
    containerImage: python:3.11-slim
    beforeCommands:
      - pip install kestra requests
    script: |
      from kestra import Kestra

      # Evaluate proxy spend metrics and health status
      primary = "{{ inputs.primary_provider }}"
      fallback = "{{ inputs.fallback_provider }}"
      budget_limit = {{ inputs.daily_budget_usd }}

      # Simulated proxy telemetry
      current_daily_spend = 124.60
      primary_healthy = True
      active_model = primary if primary_healthy else fallback
      budget_exceeded = current_daily_spend > budget_limit
      spend_utilization_pct = (current_daily_spend / budget_limit) * 100

      print(f"LiteLLM Gateway Status: Model={active_model} | Daily Spend=${current_daily_spend:.2f}/${budget_limit:.2f} ({spend_utilization_pct:.1f}%)")
      Kestra.outputs({
          "current_spend_usd": current_daily_spend,
          "spend_utilization_pct": round(spend_utilization_pct, 1),
          "active_model": active_model,
          "primary_healthy": primary_healthy,
          "budget_exceeded": budget_exceeded
      })

  - id: evaluate_governance_gate
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.check_provider_health_and_spend.vars.budget_exceeded ??
      false or not (outputs.check_provider_health_and_spend.vars.primary_healthy
      ?? true) }}"
    then:
      - id: notify_governance_alert
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Alerts AI operations and FinOps teams to failover routes or budget
          overruns.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "text": "⚠️ *LiteLLM Governance Alert:* Active route shifted or budget exceeded! Daily Spend: *${{ outputs.check_provider_health_and_spend.vars.current_spend_usd }}* ({{ outputs.check_provider_health_and_spend.vars.spend_utilization_pct }}% of limit). Active Model: `{{ outputs.check_provider_health_and_spend.vars.active_model }}` (Execution `{{ execution.id }}`)."
          }
    else:
      - id: notify_governance_nominal
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Posts routine spend digest confirmation to FinOps Slack channels.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "text": "✅ *LiteLLM Gateway Nominal:* Spend running within limits at ${{ outputs.check_provider_health_and_spend.vars.current_spend_usd }} ({{ outputs.check_provider_health_and_spend.vars.spend_utilization_pct }}% budget). Primary route `{{ inputs.primary_provider }}` healthy."
          }

  - id: archive_spend_audit_log
    type: io.kestra.plugin.scripts.shell.Commands
    containerImage: alpine:latest
    commands:
      - echo "Archived LiteLLM proxy spend metrics at $(date)"

errors:
  - id: alert_on_gateway_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Catches proxy network disconnects or API auth failures and notifies
      on-call teams.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "text": "🚨 *LiteLLM Gateway Incident:* Proxy health evaluation failed during execution `{{ execution.id }}`. Check Kestra execution logs."
      }

triggers:
  - id: hourly_spend_check
    type: io.kestra.plugin.core.trigger.Schedule
    cron: "0 * * * *"