# LiteLLM Proxy for the local governance demo.
#
# Host port 4010 -> container 4000 (4010 chosen to avoid colliding with a
# default LiteLLM on 4000). The BigQuery service-account key is mounted
# read-only from a host path kept OUTSIDE this repo (see BQ_CREDENTIALS_HOST
# in .env); it is never copied into the image or committed.

services:
  # Postgres backs LiteLLM's virtual-key / user / team management. It is what
  # enables per-user API keys, budgets, and the admin UI at /ui.
  db:
    image: postgres:16-alpine
    environment:
      POSTGRES_USER: litellm
      POSTGRES_PASSWORD: litellm
      POSTGRES_DB: litellm
    ports:
      - "${PG_HOST_PORT:-55433}:5432"   # 55433 avoids the default 5432
    volumes:
      - litellm_pgdata:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U litellm"]
      interval: 5s
      timeout: 5s
      retries: 10

  litellm:
    build:
      context: .
      dockerfile: docker/Dockerfile
    image: llm-governance-litellm:local
    ports:
      - "${PROXY_HOST_PORT:-4010}:4000"
    env_file:
      - .env
    environment:
      GOOGLE_APPLICATION_CREDENTIALS: /secrets/credentials.json
      GCP_PROJECT: "${GCP_PROJECT}"
      BQ_DATASET: "${BQ_DATASET:-llm_governance}"
      BQ_TABLE: "${BQ_TABLE:-llm_usage_logs}"
      BQ_LOCATION: "${BQ_LOCATION:-US}"
      # Virtual-key management: store keys in Postgres and enable the admin UI.
      DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
      STORE_MODEL_IN_DB: "true"
      LITELLM_SALT_KEY: "${LITELLM_SALT_KEY:-sk-salt-local-demo-change-me}"
    volumes:
      - ./litellm_config.yaml:/app/litellm_config.yaml:ro
      - ./custom_callbacks.py:/app/custom_callbacks.py:ro
      - "${BQ_CREDENTIALS_HOST}:/secrets/credentials.json:ro"
    depends_on:
      db:
        condition: service_healthy
    healthcheck:
      # The image has no curl; use python (always present in the venv).
      test: ["CMD", "python", "-c", "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://localhost:4000/health/liveliness').status==200 else 1)"]
      interval: 15s
      timeout: 5s
      retries: 5
      start_period: 40s
    restart: unless-stopped

  # Runs the budget check every 5 minutes (reuses the proxy image, which has
  # the llm_governance package + BigQuery client installed). Alerts are de-duped
  # per team/threshold/month, so a breach emails/Slacks at most once.
  scheduler:
    image: llm-governance-litellm:local
    build:
      context: .
      dockerfile: docker/Dockerfile
    env_file:
      - .env
    environment:
      GOOGLE_APPLICATION_CREDENTIALS: /secrets/credentials.json
      GCP_PROJECT: "${GCP_PROJECT}"
      BQ_DATASET: "${BQ_DATASET:-llm_governance}"
      BQ_TABLE: "${BQ_TABLE:-llm_usage_logs}"
      BUDGETS_FILE: /app/budgets.yaml
      # Reach the proxy by its compose service name (not localhost) when the
      # budget source is LiteLLM.
      LITELLM_BASE_URL: http://litellm:4000
    entrypoint: ["sh", "-c"]
    command:
      - "while true; do python -m llm_governance.cli budget-check || true; sleep 300; done"
    volumes:
      - ./budgets.yaml:/app/budgets.yaml:ro
      - "${BQ_CREDENTIALS_HOST}:/secrets/credentials.json:ro"
    depends_on:
      litellm:
        condition: service_healthy
    restart: unless-stopped

volumes:
  litellm_pgdata: