Metadata-Version: 2.4
Name: jupyterhub-litellm
Version: 0.1.0
Summary: JupyterHub LiteLLM API key management handler
License: MIT
Project-URL: Homepage, https://github.com/dive4dec/jupyterhub-litellm
Project-URL: Repository, https://github.com/dive4dec/jupyterhub-litellm
Project-URL: Issues, https://github.com/dive4dec/jupyterhub-litellm/issues
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: aiohttp
Requires-Dist: pyyaml
Requires-Dist: jupyterhub
Requires-Dist: tornado
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-asyncio; extra == "test"
Requires-Dist: aiohttp; extra == "test"

# jupyterhub-litellm

JupyterHub extension that manages per-user LiteLLM API keys — generation, budget enforcement, and admin controls. Integrates with the [LiteLLM proxy](https://docs.litellm.ai/) to give students self-service API access with configurable spending limits.

## What It Does

- Provides `/hub/litellm-key` page where each user can view and regenerate their API key
- Provides `/hub/litellm-admin` page (admin-only) for managing all users' keys and limits
- Enforces spending budgets via LiteLLM proxy's native `max_budget` + `budget_duration` rolling window — no custom spend tracking
- Shows live usage on the key page with progress bar (fetched from proxy in real time)
- Injects the API key into spawned notebooks as `LITELLM_API_KEY` / `LITELLM_API_BASE`

## Architecture

**LiteLLM proxy is the source of truth.** This extension:

- Stores only key metadata (token, expiration) in a local SQLite table (`litellm_keys`)
- Delegates all budget enforcement, rate limiting, and spend tracking to the proxy
- Queries the proxy API for live usage (`/key/info`) rather than maintaining its own counters
- No background tasks, no custom spend tables, no period-reset logic

The extension does **not** override JupyterHub's global templates or inject nav bar items directly. It only registers request handlers and ships handler-specific templates. Nav bar integration is an opt-in deployment-level concern — see [Template Customization](https://jupyterhub.readthedocs.io/en/stable/reference/template-config.html) in the JupyterHub docs.

### Handler Templates

The package ships two handler templates (`litellm_key_page.html`, `litellm-admin.html`) that extend JupyterHub's core `page.html` via `{% extends "page.html" %}`. They are findable as long as the template directory is in `c.JupyterHub.template_paths` — appended, not replacing existing paths. The package exposes the path via:

```python
from jupyterhub_litellm.utils import LITELLM_TEMPLATE_DIR
```

### Nav Bar Links

The package ships `page.html` in its templates directory, extending JupyterHub core via the `templates/` prefix (`{% extends "templates/page.html" %}`). This uses JupyterHub's `PrefixLoader` to safely extend the core without recursion, regardless of template search path order.

Nav items are injected into the left nav bar by setting `litellm_nav_items` in `template_vars`:

```python
c.JupyterHub.template_vars["litellm_nav_items"] = [
    ("LiteLLM", "litellm-key", None),
    ("LiteLLM Admin", "litellm-admin", "admin:users"),
]
```

Other packages can also override `page.html` by appending their templates **after** this package's in `template_paths` — they extend `"templates/page.html"` just as this package does, building on top of each other cooperatively via `{{ super() }}`.

## ⚠️ Critical: Master Key Synchronization

The extension authenticates with the LiteLLM proxy using `LITELLM_MASTER_KEY`. This key **must match** the proxy's configured master key.

### Common Failure Mode

If the LiteLLM proxy chart is upgraded, reinstalled, or its master key is rotated, the JupyterHub extension will fail with:

```
Failed to delete key for <user>: Authentication Error, Invalid proxy server token
```

### How to Fix

1. **Extract the current proxy master key:**
```bash
kubectl -n <litellm-namespace> get secret litellm-masterkey \
  -o go-template='{{index .data "masterkey"}}' | base64 -d
```

2. **Update your JupyterHub hub values** so `LITELLM_MASTER_KEY` matches:
```yaml
hub:
    LITELLM_MASTER_KEY: "<value-from-step-1>"
```

3. **Redeploy the hub**

### Recommended Practice

Set a **fixed** master key in your proxy values instead of relying on the auto-generated default:

```yaml
# In values/litellm/proxy.yaml
litellm-helm:
  masterkey: "sk-your-fixed-master-key"
```

This prevents the key from changing on every chart reinstall. Use the same value for `LITELLM_MASTER_KEY` in your JupyterHub config.

## How It Works

### Key Generation / Regeneration

```
Student clicks "Generate / Regenerate Key"
  │
  ├── 1. Check cooldown: was last regen within budget_duration?
  │       Yes → reject with cooldown message
  │       No  → proceed
  ├── 2. Delete old key from proxy (POST /key/delete)
  ├── 3. Build payload from current config (litellm_limits.yaml)
  ├── 4. Create new key with full budget (POST /key/generate)
  └── 5. Store metadata + last_regenerated_at in DB
```

Students can regenerate once per `budget_duration` period. Admins can regenerate at any time with no cooldown. No spend carryover — the new key gets a fresh budget, and the proxy enforces its own rolling window independently.

### Budget Enforcement

The LiteLLM proxy enforces budgets natively using a **rolling window**:

| Parameter | Type | Description |
|-----------|------|-------------|
| `max_budget` | float | Max spend in USD per `budget_duration` window |
| `budget_duration` | string | Rolling window: `"30s"`, `"15m"`, `"24h"`, `"30d"`. Spend resets automatically as old usage ages out. |

When `budget_duration` is `"24h"`, the proxy enforces: *"In the last 24 hours, has this key spent more than max_budget?"* — no manual resets needed.

### Key Regeneration — Cooldown

Students can regenerate their API key **once per budget_duration period**. This prevents abuse while giving students flexibility to recover from accidental leaks.

- **Student**: cooldown enforced — must wait one full `budget_duration` before regenerating again
- **Admin**: no cooldown — can regenerate any user's key at any time
- **No spend carryover**: regenerated keys get the full budget. The proxy enforces its rolling window independently.

If a student tries to regenerate before the cooldown expires, they see:
> You can regenerate your key once per budget duration period. Try again in X.X hours.

The last regeneration time is visible on the admin page for reference.

## Configuration

Budget limits are defined in `jupyterhub_litellm/litellm_limits.yaml`:

```yaml
defaults:
  max_budget: 100            # $100 per budget_duration window
  budget_duration: "24h"     # resets every 24 hours (rolling window)
  tpm_limit: null            # tokens per minute (null = no limit)
  rpm_limit: null            # requests per minute (null = no limit)
  duration: "30d"            # key lifetime before auto-deactivation

per_user:
  # Uncomment to override for specific users
  # username:
  #   max_budget: null        # no budget limit
  #   budget_duration: null   # no auto-reset
  #   tpm_limit: null
  #   rpm_limit: null
  #   duration: "30d"
```

### Field Reference

All fields map directly to [LiteLLM proxy virtual key parameters](https://docs.litellm.ai/docs/proxy/virtual_keys):

| Field | Type | Description | Default |
|-------|------|-------------|---------|
| `max_budget` | float / null | Max spend in USD per `budget_duration` window. Set to `null` for unlimited. | `null` (unlimited) |
| `budget_duration` | string / null | Rolling window: `"30s"`, `"15m"`, `"24h"`, `"30d"`. Set to `null` for no auto-reset. | `null` |
| `tpm_limit` | int / null | Tokens per minute limit. Set to `null` for no limit. | `null` |
| `rpm_limit` | int / null | Requests per minute limit. Set to `null` for no limit. | `null` |
| `duration` | string | Key lifetime after which key deactivates: `"30d"`, `"24h"`, etc. | `"30d"` |

## UI

### Key Page (`/hub/litellm-key`)

- **API Key** — copyable key with expiration date
- **Usage bar** — `$spent / $max_budget` with progress percentage (green/yellow/red)
- **Limits** — displays `budget_duration`, `tpm_limit`, `rpm_limit` if set
- **Generate / Regenerate Key** button

### Admin Page (`/hub/litellm-admin`)

| Column | Description |
|--------|-------------|
| User | Username |
| Spent | Current proxy spend ($XX.XXXX) |
| Max Budget | Budget limit per window |
| Budget Duration | Rolling window (e.g. 24h) |
| Last Regen | Last regeneration timestamp (admin only) |
| TPM/RPM | Rate limits |
| Blocked | Whether the key is blocked |
| Expires | Key expiration date |

**Actions per user:**
- ⟳ **Regenerate** — create new key, carrying over current spend
- 🔒 **Block** — immediately block the key (rejects all requests)
- 🔓 **Unblock** — unblock a previously blocked key
- ⚙ **Edit Limits** — modify `max_budget`, `budget_duration`, `tpm_limit`, `rpm_limit`

## Deployment

Add the following to your `jupyterhub_config.py` (or equivalent Helm `extraConfig`):

```python
import sys, os

# Ensure the package is importable — adjust the path to wherever it lives
pkg_dir = "/path/to/jupyterhub-litellm"
sys.path.insert(0, pkg_dir)

# Append the handler template directory (required for handler templates to resolve)
from jupyterhub_litellm.utils import LITELLM_TEMPLATE_DIR
c.JupyterHub.template_paths.append(LITELLM_TEMPLATE_DIR)

# Register the handlers
import jupyterhub_litellm
from jupyterhub_litellm.admin import LiteLLMAdminHandler
c.JupyterHub.extra_handlers.append((r"litellm-key", jupyterhub_litellm.LiteLLMKeyHandler))
c.JupyterHub.extra_handlers.append((r"litellm-admin", LiteLLMAdminHandler))
```

### Environment Variables

Set these so the extension can reach your LiteLLM proxy and budget config:

| Variable | Description | Example |
|----------|-------------|---------|
| `LITELLM_MASTER_KEY` | Proxy admin key for API calls. **Must match proxy's master key.** See [Master Key Synchronization](#critical-master-key-synchronization). | `sk-...` |
| `LITELLM_PROXY_URL` | Proxy endpoint | `http://litellm.litellm-proxy.svc.cluster.local:4000` |
| `LITELLM_DB_PATH` | SQLite DB for key metadata | `/srv/jupyterhub/jupyterhub.sqlite` |
| `LITELLM_API_ENDPOINT` | External API URL shown to users | `https://example.com/litellm/v1` |
| `LITELLM_MODEL_NAME` | Model name shown to users | `Socrates` |
| `LITELLM_LIMITS_PATH` | Path to `litellm_limits.yaml` | `/path/to/jupyterhub_litellm/litellm_limits.yaml` |

## Database Schema

One table in the JupyterHub SQLite database:

```sql
-- Per-user API key metadata (no spend tracking — proxy is source of truth)
CREATE TABLE litellm_keys (
    username          TEXT PRIMARY KEY,
    key               TEXT NOT NULL,
    key_name          TEXT,
    token_id          TEXT,
    created_at        TEXT,
    expires_at        TEXT,
    last_regenerated_at TEXT DEFAULT NULL  -- ISO timestamp, for student regen cooldown
);
```

## Known Issues

- **Stale `__pycache__`**: K8s `extraFiles` sets epoch (1970) timestamps on extracted files. Fixed by `PYTHONDONTWRITEBYTECODE=1` + startup `__pycache__` cleanup in `extraConfig`.
- **Master key mismatch after proxy upgrade**: If the proxy chart auto-generates a new master key, `LITELLM_MASTER_KEY` in JupyterHub must be updated. Use a fixed `masterkey` in proxy values to prevent this.
