commit 4c9ed3c24b909d60139aa003dee8d8807647f4f9 Author: Daniel Arroyo Date: Thu Jul 30 10:58:55 2026 -0400 Initial commit: LlamaLink Go rewrite Complete rewrite from Python/FastAPI to Go/Gin: - Go backend: auth (API keys + bcrypt), llama.cpp subprocess manager, hot-swap multi-model, rate limiting, quota system, webhooks - Vue 3 SPA admin panel (src/) with Tailwind CSS - Deployment: Docker multi-stage, docker-compose, nginx, systemd - GORM/SQLite models: ApiKey, Model, UsageLog, Quota, Webhook - REST API: /api/v1/admin/* (keys, models, chat, usage, health) - Embedded frontend via go:embed (build output at web/dist/) Removed legacy Python artifacts (app/, tests/, pyproject.toml, etc.) diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..6dacce2 --- /dev/null +++ b/.env.example @@ -0,0 +1,28 @@ +LLAMALINK_ENV=development +LLAMALINK_HOST=0.0.0.0 +LLAMALINK_PORT=8000 +LLAMALINK_SECRET_KEY=change-me-in-production + +DATABASE_URL=sqlite:///./llamalink.db +DATABASE_MAX_OPENConns=25 +DATABASE_MAX_IDLE_CONNS=5 +DATABASE_CONN_MAX_LIFETIME=300 + +MANAGE_LLAMA_SERVER=true +LLAMA_SERVER_BIN=/usr/local/bin/llama-server +LLAMA_SERVER_HOST=127.0.0.1 +LLAMA_SERVER_PORT=8080 +LLAMA_SERVER_STARTUP_TIMEOUT=120 +LLAMA_SERVER_STOP_TIMEOUT=10 +MODEL_SWAP_COOLDOWN=2 + +RATE_LIMIT_PER_MINUTE=60 +RATE_LIMIT_STORAGE=memory +# RATE_LIMIT_REDIS_URL=redis://localhost:6379/0 + +ADMIN_TOKEN=change-me-in-production + +LOG_LEVEL=info +LOG_FORMAT=json + +# WEBHOOK_SECRET=your-webhook-secret-here diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e57d6bb --- /dev/null +++ b/.gitignore @@ -0,0 +1,75 @@ +# Go binaries +/llamalink +/llamalink-* +/dist/ +/bin/ + +# Test artifacts +*.test +*.out +coverage.html +coverage.out + +# Vendoring — using go modules +/vendor/ + +# Go workspace +go.work +go.work.sum + +# Environment +.env +.env.local +.env.*.local +*.env +!.env.example + +# Database +*.db +*.db-shm +*.db-wal +*.sqlite +*.sqlite3 +/data/ + +# LLM models y binarios +/models/ +*.gguf +/llama-server +/llama-cli + +# Logs +*.log +logs/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Frontend +web/frontend/node_modules/ +web/frontend/dist/ +web/frontend/.vite/ +web/dist/ +web/frontend/.env +web/frontend/.env.local +web/frontend/coverage/ + +# Build artifacts +*.tmp +*.bak +tmp/ +.cache/ + +# Coverage reports +*.coverprofile + +# Misc +vendor/ diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..99b1275 --- /dev/null +++ b/Makefile @@ -0,0 +1,104 @@ +.PHONY: build run test lint clean migrate dev prod deps fmt + +# Binary name +BINARY=llamalink +VERSION=$(shell git describe --tags --always --dirty 2>/dev/null || echo "dev") +BUILD_TIME=$(shell date -u '+%Y-%m-%d_%H:%M:%S') +LDFLAGS=-ldflags "-X main.Version=$(VERSION) -X main.BuildTime=$(BUILD_TIME)" + +# Directories +BUILD_DIR=./bin +FRONTEND_DIR=./web/frontend + +# Go parameters +GOCMD=go +GOBUILD=$(GOCMD) build $(LDFLAGS) +GOTEST=$(GOCMD) test +GOGET=$(GOCMD) get +GOMOD=$(GOCMD) mod +GOFMT=gofmt +GOLINT=golangci-lint + +# Default target +all: deps build + +## build: Build the binary +build: + @echo "Building $(BINARY)..." + @mkdir -p $(BUILD_DIR) + $(GOBUILD) -o $(BUILD_DIR)/$(BINARY) ./cmd/llamalink + +## run: Build and run +run: build + @echo "Running..." + $(BUILD_DIR)/$(BINARY) + +## dev: Run in development mode +dev: + LLAMALINK_ENV=development $(GOCMD) run ./cmd/llamalink + +## test: Run tests +test: + $(GOTEST) -v -race -coverprofile=coverage.out ./... + +## test/integration: Run integration tests +test/integration: + $(GOTEST) -v -tags=integration ./tests/integration/... + +## lint: Run linters +lint: + $(GOLINT) run ./... + +## fmt: Format code +fmt: + $(GOFMT) -s -w . + +## deps: Download dependencies +deps: + $(GOMOD) download + $(GOMOD) tidy + +## clean: Remove build artifacts +clean: + rm -rf $(BUILD_DIR) + rm -f coverage.out + rm -f *.db + +## migrate/up: Run database migrations up +migrate/up: + migrate -path internal/db/migrations -database "$(DATABASE_URL)" up + +## migrate/down: Run database migrations down +migrate/down: + migrate -path internal/db/migrations -database "$(DATABASE_URL)" down + +## migrate/create: Create a new migration +migrate/create NAME=add_users_table: + migrate create -path internal/db/migrations -ext .sql -dir internal/db/migrations $(NAME) + +## frontend/install: Install frontend dependencies +frontend/install: + cd $(FRONTEND_DIR) && npm install + +## frontend/dev: Run frontend dev server +frontend/dev: + cd $(FRONTEND_DIR) && npm run dev + +## frontend/build: Build frontend for production +frontend/build: + cd $(FRONTEND_DIR) && npm run build + +## docker/build: Build Docker image +docker/build: + docker build -t llamalink:latest -f deploy/Dockerfile . + +## docker/run: Run Docker container +docker/run: + docker-compose -f deploy/docker-compose.yml up + +## docker/build/run: Build and run with docker-compose +docker/build/run: docker/build docker/run + +## help: Show this help +help: + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' diff --git a/README.md b/README.md new file mode 100644 index 0000000..d3671ce --- /dev/null +++ b/README.md @@ -0,0 +1,189 @@ +# LlamaLink + +**API Gateway en Go** para ejecutar modelos de lenguaje (LLMs) de forma local usando **llama.cpp**, con autenticación por API keys, rate limiting, cuotas de uso y panel de administración Vue 3. + +## Características + +- **API OpenAI-compatible**: `POST /v1/chat/completions` con streaming SSE +- **Multi-model hot-swap**: Configura muchos modelos, activa uno a la vez sin downtime +- **Auth por API keys**: Keys con scopes (`chat`, `models`, `usage`, `admin`), hashing bcrypt +- **Rate limiting**: Token bucket por API key (configurable, backend memory o Redis) +- **Cuotas mensuales**: Tracking por modelo o global, webhooks al 90% y 100% +- **Admin SPA**: Vue 3 + Tailwind CSS con Dashboard, Keys, Models, Usage +- **Deployment**: Docker, docker-compose, systemd, nginx/Caddy +- **WebSocket**: Live updates de estado del modelo y stats + +## Requisitos + +- Go 1.23+ +- Node.js 20+ (para build del frontend) +- llama-server (GGML/GGUF) +- CUDA 12.x (opcional, para GPU) + +## Quick Start + +### 1. Build + +```bash +# Dependencias Go +go mod download + +# Build binario +make build + +# O directamente +go build -o llamalink ./cmd/llamalink +``` + +### 2. Configuración + +```bash +cp .env.example .env +# Editar .env - mínimo: ADMIN_TOKEN +``` + +### 3. Arrancar + +```bash +# Con llama-server externo (ya corriendo en 127.0.0.1:8080) +./llamalink + +# O gestionar llama-server internamente +MANAGE_LLAMA_SERVER=true ./llamalink +``` + +### 4. Probar + +```bash +# Health check +curl http://localhost:8000/health + +# Login admin (primera key se crea automáticamente) +ADMIN_TOKEN=tu-token curl http://localhost:8000/api/v1/admin/login \ + -X POST -H "Content-Type: application/json" \ + -d '{"admin_token":"tu-token"}' + +# Crear API key +curl http://localhost:8000/v1/keys \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"name":"test-key"}' + +# Chat completion +curl http://localhost:8000/v1/chat/completions \ + -H "Authorization: Bearer $TU_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"tu-modelo","messages":[{"role":"user","content":"Hello!"}]}' +``` + +## Variables de Entorno + +| Variable | Descripción | Default | +|---|---|---| +| `LLAMALINK_ENV` | `development` o `production` | `development` | +| `LLAMALINK_PORT` | Puerto HTTP | `8000` | +| `DATABASE_URL` | Connection string SQLite/Postgres | `sqlite:///./llamalink.db` | +| `MANAGE_LLAMA_SERVER` | ¿Gestionar llama-server internamente? | `true` | +| `LLAMA_SERVER_BIN` | Path a `llama-server` | `/usr/local/bin/llama-server` | +| `LLAMA_SERVER_HOST` | Host de llama-server | `127.0.0.1` | +| `LLAMA_SERVER_PORT` | Puerto de llama-server | `8080` | +| `LLAMA_SERVER_STARTUP_TIMEOUT` | Timeout startup (seg) | `120` | +| `RATE_LIMIT_PER_MINUTE` | Requests/min por key | `60` | +| `ADMIN_TOKEN` | Token admin inicial | `changeme` | +| `LOG_LEVEL` | `debug`, `info`, `warn`, `error` | `info` | + +## Docker + +```bash +# Build imagen +make docker/build + +# O docker-compose completo (API + llama-runner + nginx) +docker compose -f deploy/docker-compose.yml up -d + +# Production con tu modelo +ADMIN_TOKEN=mi-token-secreto docker compose -f deploy/docker-compose.yml up -d +``` + +## Admin Panel + +Accede a `http://localhost:8000/admin/` para el panel Vue 3: + +- **Dashboard**: Stats, estado del modelo, tabla de modelos con hot-swap +- **API Keys**: Crear, listar, revocar keys +- **Models**: CRUD de modelos, load/unload +- **Usage**: Charts de uso, logs, quota + +## API Endpoints + +### OpenAI-compatible +- `POST /v1/chat/completions` — Chat con streaming +- `GET /v1/models` — Lista modelos +- `GET /v1/models/active` — Modelo activo + +### Gestión +- `GET|POST /v1/keys` — Listar/crear keys (admin) +- `DELETE /v1/keys/{id}` — Revocar key (admin) +- `POST /v1/models` — Crear modelo (admin) +- `POST /v1/models/{name}/load` — Hot-swap (admin) +- `GET /v1/usage` — Uso de tu key +- `GET /v1/usage/{key_id}` — Uso específico (admin) + +### Admin API (JSON) +- `POST /api/v1/admin/login` — Login admin +- `GET /api/v1/admin/dashboard` — Stats aggregated +- `GET /api/v1/admin/keys` — Keys con detalles +- `POST /api/v1/admin/keys` — Crear key +- `DELETE /api/v1/admin/keys/{id}` — Revocar + +## Makefile + +```bash +make build # Build binario +make run # Build y ejecutar +make dev # go run (desarrollo) +make test # Tests +make lint # golangci-lint +make clean # Limpiar build +make migrate/up # Correr migrations +make frontend/build # Build Vue SPA +``` + +## Arquitectura + +``` +llamalink/ +├── cmd/llamalink/ # Entry point +├── internal/ +│ ├── api/ # Gin router, handlers, middleware +│ ├── auth/ # API key service, bcrypt +│ ├── db/ # GORM models, migrations +│ ├── llama/ # Proxy, subprocess manager +│ └── quota/ # Quota check, webhooks +├── web/frontend/ # Vue 3 SPA (build → web/dist/) +├── deploy/ # Docker, systemd, nginx +└── migrations/ # SQL migrations +``` + +## Deployment + +### Systemd + +```bash +sudo cp deploy/llamalink.service /etc/systemd/system/ +sudo systemctl daemon-reload +sudo systemctl enable llamalink +sudo systemctl start llamalink +``` + +###nginx + +```bash +sudo cp deploy/nginx.conf /etc/nginx/sites-available/llamalink +sudo ln -s /etc/nginx/sites-available/llamalink /etc/nginx/sites-enabled/ +sudo nginx -t && sudo systemctl reload nginx +``` + +## License + +MIT diff --git a/SPEC.md b/SPEC.md new file mode 100644 index 0000000..1ca4661 --- /dev/null +++ b/SPEC.md @@ -0,0 +1,285 @@ +# LlamaLink — Especificación Técnica (Living Document) + +> Este documento se actualiza conforme evoluciona la implementación. Última actualización: implementación inicial. + +## 1. Descripción General + +**LlamaLink** es un servidor middleware/gateway que permite ejecutar modelos de lenguaje (LLMs) de forma local usando **llama.cpp**, exponiéndolos de manera segura a través de una API REST propia protegida con **API keys**, permitiendo así el acceso remoto controlado a la inferencia. + +### Objetivos principales +- Ejecutar modelos GGUF localmente mediante llama.cpp. +- Exponer un endpoint HTTP compatible con el estándar OpenAI (`/v1/chat/completions`). +- Autenticar y autorizar clientes remotos mediante API keys. +- Controlar uso, límites y seguridad del acceso a los modelos. + +--- + +## 2. Arquitectura + +``` +Cliente remoto (con API key) + │ + ▼ +┌─────────────────────┐ +│ LlamaLink API │ ← Gateway (FastAPI) +│ - Autenticación │ +│ - Rate limiting │ +│ - Enrutamiento │ +│ - Logs y métricas │ +└──────────┬───────────┘ + │ (proxy HTTP) + ▼ +┌─────────────────────┐ +│ llama.cpp │ ← Motor de inferencia +│ (llama-server) │ +└─────────────────────┘ +``` + +### Integración con llama.cpp + +- **Opción A** (implementada): LlamaLink actúa como gateway/proxy delante del servidor HTTP nativo de llama.cpp (`llama-server`). +- **Opción B**: Uso de `llama-cpp-python` embebido (futuro). + +### Multi-modelo + +- Se soporta la configuración de múltiples modelos en la base de datos. +- Solo **un modelo** está activo/cargado en llama-server en un momento dado. +- Si llega una request para un modelo diferente al activo: + 1. Se rechaza con HTTP 503 + `Retry-After` mientras dura el swap + 2. El admin puede disparar el swap manualmente vía `POST /v1/models/{name}/load` + 3. Ocurre automáticamente en lazy mode si `MANAGE_LLAMA_SERVER=true` +- El swap es secuencial (un `asyncio.Lock` evita swaps concurrentes). + +--- + +## 3. Stack Tecnológico + +| Componente | Tecnología | Justificación | +|---|---|---| +| Motor de inferencia | **llama.cpp** (C++) | Ya compilado y optimizado, no requiere reescritura | +| API Gateway | **Python 3.11+ / FastAPI** | Desarrollo rápido, ecosistema maduro, buena documentación automática (OpenAPI/Swagger) | +| Servidor ASGI | **Uvicorn** | Estándar para FastAPI, soporta async y alto rendimiento | +| Base de datos | **SQLite** (MVP) → **PostgreSQL** (producción) | Escalable según necesidad | +| ORM | **SQLAlchemy 2.0** (async) | Manejo de modelos y migraciones | +| Autenticación | API Keys con hash **SHA-256** | Nunca se almacenan en texto plano | +| Rate limiting | **slowapi** (memoria o Redis) | Control de abuso por IP | +| Reverse proxy / TLS | **Nginx** o **Caddy** | HTTPS obligatorio en producción | +| Contenedores | **Docker + docker-compose** | Despliegue reproducible | + +--- + +## 4. Estructura de Carpetas + +``` +llamalink/ +├── app/ +│ ├── main.py # Punto de entrada FastAPI + lifespan +│ ├── config.py # Configuración pydantic-settings +│ ├── __init__.py +│ ├── core/ +│ │ ├── __init__.py +│ │ ├── security.py # Hash/verify API keys (SHA-256) +│ │ ├── errors.py # JSON error handlers estandarizados +│ │ └── logging.py # Structured logging +│ ├── db/ +│ │ ├── __init__.py +│ │ ├── base.py # SQLAlchemy declarative base +│ │ ├── session.py # Engine + session factory +│ │ └── models.py # ORM models +│ ├── schemas/ +│ │ ├── __init__.py +│ │ ├── common.py # ErrorResponse, HealthResponse +│ │ ├── chat.py # OpenAI-compat request/response +│ │ └── keys.py # Key management schemas +│ ├── auth/ +│ │ ├── __init__.py +│ │ ├── service.py # CRUD keys, hashing +│ │ └── dependencies.py # Bearer auth FastAPI deps +│ ├── routers/ +│ │ ├── __init__.py +│ │ ├── chat.py # /v1/chat/completions +│ │ ├── keys.py # /v1/keys CRUD +│ │ └── health.py # /health, /ready +│ └── services/ +│ ├── __init__.py +│ ├── llama_proxy.py # httpx async client → llama-server +│ ├── rate_limiter.py # slowapi setup +│ ├── usage_tracker.py # Token counting + DB logs +│ ├── quota.py # Chequeo de cuota mensual +│ ├── webhook.py # Notificaciones al superar cuota +│ ├── model_manager.py # Subprocess lifecycle + state machine +│ └── models_registry.py # Mapa de modelos configurados +├── migrations/ # Alembic +├── deploy/ # Nginx, Caddy, systemd units +├── docker/ # Dockerfiles +├── scripts/ # CLI tools +├── tests/ +├── pyproject.toml +├── .env.example +└── README.md +``` + +--- + +## 5. Endpoints Implementados + +| Método | Endpoint | Descripción | Auth | +|---|---|---|---| +| POST | `/v1/chat/completions` | Enviar prompt y recibir respuesta del modelo | Sí | +| POST | `/v1/keys` | Crear nueva API key | Sí (admin) | +| GET | `/v1/keys` | Listar API keys existentes | Sí (admin) | +| DELETE | `/v1/keys/{key_id}` | Revocar una API key | Sí (admin) | +| GET | `/health` | Verificar estado del servicio | No | +| GET | `/ready` | Verificar listo para servir | No | + +### Pendiente de implementar (Fase 4+) + +| Método | Endpoint | Descripción | +|---|---|---| +| GET | `/v1/models` | Listar modelos configurados | +| GET | `/v1/models/active` | Mostrar cuál está cargado | +| POST | `/v1/models/{name}/load` | Cargar/swap a un modelo | +| GET | `/v1/usage/{key_id}` | Métricas de uso de una key | +| GET/DELETE | `/admin/*` | Panel web de administración | + +--- + +## 6. Variables de Entorno + +```env +LLAMALINK_ENV=development +LLAMALINK_HOST=0.0.0.0 +LLAMALINK_PORT=8000 +LLAMALINK_SECRET_KEY=changeme + +DATABASE_URL=sqlite+aiosqlite:///./llamalink.db + +MANAGE_LLAMA_SERVER=true +LLAMA_SERVER_BIN=llama-server +LLAMA_SERVER_HOST=127.0.0.1 +LLAMA_SERVER_PORT=8080 +LLAMA_SERVER_STARTUP_TIMEOUT=120 +MODEL_SWAP_COOLDOWN=2 +LLAMA_SERVER_STOP_TIMEOUT=10 + +RATE_LIMIT_PER_MINUTE=60 +RATE_LIMIT_STORAGE=memory + +LOG_LEVEL=INFO +LOG_FORMAT=json + +ADMIN_API_KEY=changeme +``` + +--- + +## 7. Modelo de Datos + +### ApiKey +- `id`, `name`, `hashed_key`, `key_prefix` +- `scopes`: CSV de permisos (`chat,models,usage,admin`) +- `is_active`, `is_admin`, `created_at`, `last_used_at`, `owner_label` + +### Model +- `id`, `name`, `model_path`, `alias` +- `ctx_size`, `n_gpu_layers`, `extra_args` (JSON) +- `is_default`, `is_enabled`, `is_active`, `loaded_at` + +### UsageLog +- `api_key_id`, `model_name`, `endpoint` +- `prompt_tokens`, `completion_tokens`, `total_tokens` +- `latency_ms`, `status`, `ip_address`, `user_agent` +- `error_message`, `created_at` + +### Quota +- `api_key_id`, `model_scope` (None = global) +- `period_start`, `period_end` +- `tokens_limit`, `tokens_used` + +### Webhook +- `api_key_id`, `url`, `event`, `secret` (HMAC) +- `is_active`, `created_at` + +--- + +## 8. Estados del ModelManager + +``` +STOPPED ──load()──> LOADING ──success──> READY + ↑ │ + │ fail + │ ↓ + └──<─────── FAILED FAILED + │ + └──swap()──> SWAPPING ──success──> READY +``` + +--- + +## 9. Formato de Errores + +Todos los errores siguen este formato: + +```json +{ + "error": { + "code": "error_code_string", + "message": "Descripción legible", + "retry_after_seconds": 30 + } +} +``` + +Códigos: `invalid_api_key`, `api_key_revoked`, `rate_limit_exceeded`, `model_not_found`, `model_loading`, `quota_exceeded`, `upstream_error`, `internal_error`, `validation_error`, `unauthorized` + +--- + +## 10. Roadmap de Implementación + +### ✅ Fase 0 — Esqueleto +- pyproject.toml, config, main, db, .env.example, .gitignore + +### ✅ Fase 1 — MVP +- Auth (hash, service, deps) +- llama_proxy (httpx async, streaming) +- /v1/chat/completions (non-stream + stream) +- /v1/keys CRUD +- /health, /ready +- Error handlers estandarizados +- Quota service (stub) +- Rate limiter (stub) + +### 🔄 Fase 2 — Endurecimiento +- [ ] Rate limiting (slowapi) +- [ ] Usage tracker middleware +- [ ] Logging estructurado JSON + +### 📋 Fase 3 — Docker + Producción +- [ ] docker-compose.yml +- [ ] Dockerfile.api, Dockerfile.llama +- [ ] Nginx.conf, Caddyfile +- [ ] systemd units + +### 📋 Fase 4 — Multi-modelo + Streaming +- [ ] model_manager (subprocess lifecycle, state machine) +- [ ] streaming en /v1/chat/completions +- [ ] /v1/models CRUD +- [ ] Lazy swap + admin trigger +- [ ] 503 + Retry-After durante swap + +### 📋 Fase 5 — Cuotas + Webhooks +- [ ] Quota enforcement middleware +- [ ] Webhook delivery con HMAC +- [ ] Background task para notificaciones + +### 📋 Fase 6 — Panel Admin Web +- [ ] Jinja2 templates + htmx +- [ ] /admin/* routes +- [ ] Crear/revocar keys desde UI +- [ ] Ver métricas y logs + +### 📋 Fase 7 — Polish +- [ ] Validación estricta de inputs +- [ ] Auto-revoke en actividad sospechosa +- [ ] Cobertura de tests ≥ 80% +- [ ] ruff + mypy limpios diff --git a/cmd/llamalink/main.go b/cmd/llamalink/main.go new file mode 100644 index 0000000..bbd49fa --- /dev/null +++ b/cmd/llamalink/main.go @@ -0,0 +1,107 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "github.com/llamalink/llamalink/internal/api" + "github.com/llamalink/llamalink/internal/config" + "github.com/llamalink/llamalink/internal/db" + "github.com/llamalink/llamalink/internal/llama" +) + +var ( + Version = "dev" + BuildTime = "unknown" +) + +func main() { + cfg := config.Load() + logger := cfg.Logger() + slog.SetDefault(logger) + + logger.Info("starting llamalink", + "version", Version, + "build_time", BuildTime, + "env", cfg.LlamalinkEnv, + ) + + // Initialize database + database, err := db.Open(cfg) + if err != nil { + logger.Error("failed to open database", "error", err) + os.Exit(1) + } + defer func() { + sqlDB, _ := database.DB() + if sqlDB != nil { + sqlDB.Close() + } + }() + + // Run migrations + if err := db.Migrate(database); err != nil { + logger.Error("failed to run migrations", "error", err) + os.Exit(1) + } + logger.Info("database migrations complete") + + // Seed admin key if needed + if err := db.SeedAdminKey(database, cfg.AdminToken); err != nil { + logger.Error("failed to seed admin key", "error", err) + os.Exit(1) + } + + // Initialize llama manager + var llamaManager *llama.Manager + if cfg.ManageLlamaServer { + llamaManager = llama.NewManager(cfg, database) + if err := llamaManager.Start(); err != nil { + logger.Warn("failed to start llama manager", "error", err) + } + defer llamaManager.Stop() + } + + // Setup router + router := api.New(cfg, database, llamaManager) + + // Server + addr := fmt.Sprintf("%s:%d", cfg.LlamalinkHost, cfg.LlamalinkPort) + srv := &http.Server{ + Addr: addr, + Handler: router, + ReadTimeout: 30 * time.Second, + WriteTimeout: 60 * time.Second, + IdleTimeout: 120 * time.Second, + } + + // Graceful shutdown + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + + go func() { + logger.Info("server listening", "addr", addr) + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + logger.Error("server error", "error", err) + os.Exit(1) + } + }() + + <-quit + logger.Info("shutting down server...") + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if err := srv.Shutdown(ctx); err != nil { + logger.Error("server forced to shutdown", "error", err) + } + + logger.Info("server stopped") +} diff --git a/deploy/Caddyfile b/deploy/Caddyfile new file mode 100644 index 0000000..8d87a65 --- /dev/null +++ b/deploy/Caddyfile @@ -0,0 +1,60 @@ +# LlamaLink Caddyfile +# Auto HTTPS via Let's Encrypt + +llamalink.local { + reverse_proxy llamalink:8000 + + log { + output file /var/log/caddy/llamalink.log + } + + handle /health { + reverse_proxy llamalink:8000 + } + + handle /ready { + reverse_proxy llamalink:8000 + } + + handle /v1/chat/completions { + reverse_proxy llamalink:8000 { + flush_interval -1 + } + } + + handle { + reverse_proxy llamalink:8000 + } +} + +# Production with TLS +# Replace with your domain +llamalink.example.com { + reverse_proxy llamalink:8000 + + tls { + protocols tls1.2 tls1.3 + } + + log { + output file /var/log/caddy/llamalink.log + } + + handle /health { + reverse_proxy llamalink:8000 + } + + handle /ready { + reverse_proxy llamalink:8000 + } + + handle /v1/chat/completions { + reverse_proxy llamalink:8000 { + flush_interval -1 + } + } + + handle { + reverse_proxy llamalink:8000 + } +} diff --git a/deploy/Dockerfile b/deploy/Dockerfile new file mode 100644 index 0000000..821d085 --- /dev/null +++ b/deploy/Dockerfile @@ -0,0 +1,53 @@ +# Build stage +FROM node:20-alpine AS node-builder + +WORKDIR /app + +COPY web/frontend/package*.json ./ +RUN npm ci + +COPY web/frontend/ ./ +RUN npm run build + +# Go stage +FROM golang:1.23-alpine AS go-builder + +RUN apk add --no-cache git ca-certificates + +WORKDIR /app + +COPY go.mod go.sum ./ +RUN go mod download + +COPY . . + +# Copy frontend build +COPY --from=node-builder /app/dist ./web/dist + +RUN CGO_ENABLED=0 GOOS=linux go build -ldflags="-w -s" -o llamalink ./cmd/llamalink + +# Final stage +FROM alpine:3.19 + +RUN apk add --no-cache ca-certificates curl + +WORKDIR /app + +# Create non-root user +RUN addgroup -g 1000 llamalink && \ + adduser -u 1000 -G llamalink -s /bin/sh -D llamalink + +COPY --from=go-builder /app/llamalink . +COPY --from=go-builder /app/.env.example .env + +# Create data directory +RUN mkdir -p /app/data && chown llamalink:llamalink /app/data + +USER llamalink + +EXPOSE 8000 + +ENV LLAMALINK_HOST=0.0.0.0 +ENV LLAMALINK_PORT=8000 + +ENTRYPOINT ["./llamalink"] diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml new file mode 100644 index 0000000..5a9445e --- /dev/null +++ b/deploy/docker-compose.yml @@ -0,0 +1,79 @@ +services: + llamalink: + build: + context: . + dockerfile: deploy/Dockerfile + container_name: llamalink-api + ports: + - "8000:8000" + environment: + - LLAMALINK_ENV=production + - DATABASE_URL=sqlite+aiosqlite:///./data/llamalink.db + - MANAGE_LLAMA_SERVER=true + - LLAMA_SERVER_HOST=127.0.0.1 + - LLAMA_SERVER_PORT=8080 + - LLAMA_SERVER_BIN=/usr/local/bin/llama-server + - LLAMA_SERVER_STARTUP_TIMEOUT=120 + - MODEL_SWAP_COOLDOWN=2 + - LLAMA_SERVER_STOP_TIMEOUT=10 + - RATE_LIMIT_PER_MINUTE=60 + - RATE_LIMIT_STORAGE=memory + - LOG_LEVEL=info + - LOG_FORMAT=json + - ADMIN_TOKEN=${ADMIN_TOKEN} + volumes: + - llamalink-data:/app/data + - ./models:/models:ro + restart: unless-stopped + networks: + - llamalink-net + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:8000/health"] + interval: 30s + timeout: 10s + retries: 3 + start_period: 10s + + llama-runner: + image: ghcr.io/ggml-org/llama.cpp:server + container_name: llama-runner + environment: + - CUDA_VISIBLE_DEVICES=0 + volumes: + - ./models:/models:ro + restart: unless-stopped + networks: + - llamalink-net + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:8080/health"] + interval: 30s + timeout: 10s + retries: 3 + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + + nginx: + image: nginx:1.27-alpine + container_name: llamalink-nginx + ports: + - "80:80" + - "443:443" + volumes: + - ./deploy/nginx.conf:/etc/nginx/nginx.conf:ro + depends_on: + - llamalink + restart: unless-stopped + networks: + - llamalink-net + +volumes: + llamalink-data: + +networks: + llamalink-net: + driver: bridge diff --git a/deploy/llama-server@.service b/deploy/llama-server@.service new file mode 100644 index 0000000..b303950 --- /dev/null +++ b/deploy/llama-server@.service @@ -0,0 +1,35 @@ +[Unit] +Description=llama-server instance %i +After=network.target + +[Service] +Type=simple +User=llamalink +WorkingDirectory=/opt/llamalink +ExecStart=/usr/local/bin/llama-server \ + --model /opt/llamalink/models/%i.gguf \ + --alias %i \ + --host 127.0.0.1 \ + --port 8080 \ + --ctx-size 8192 \ + --n-gpu-layers auto \ + --parallel 4 \ + --rope-scaling linear +Restart=on-failure +RestartSec=5 +StandardOutput=journal +StandardError=journal +SyslogIdentifier=llama-server-%i + +Environment="CUDA_VISIBLE_DEVICES=0" + +# Security hardening +NoNewPrivileges=true +PrivateTmp=true +ProtectSystem=strict +ProtectHome=true +ReadOnlyPaths=/opt/llamalink/models +ReadWritePaths=/opt/llamalink/data + +[Install] +WantedBy=multi-user.target diff --git a/deploy/llamalink.service b/deploy/llamalink.service new file mode 100644 index 0000000..bd70466 --- /dev/null +++ b/deploy/llamalink.service @@ -0,0 +1,28 @@ +[Unit] +Description=LlamaLink API Gateway +After=network.target + +[Service] +Type=simple +User=llamalink +WorkingDirectory=/opt/llamalink +ExecStart=/usr/local/bin/llamalink \ + --host 0.0.0.0 \ + --port 8000 +Restart=on-failure +RestartSec=5 +StandardOutput=journal +StandardError=journal +SyslogIdentifier=llamalink + +# Security +NoNewPrivileges=true +PrivateTmp=true +ProtectSystem=strict +ProtectHome=true +ReadWritePaths=/opt/llamalink/data +ReadOnlyPaths=/opt/llamalink/models +Environment=LLAMALINK_ENV=production + +[Install] +WantedBy=multi-user.target diff --git a/deploy/nginx.conf b/deploy/nginx.conf new file mode 100644 index 0000000..3b52dc7 --- /dev/null +++ b/deploy/nginx.conf @@ -0,0 +1,122 @@ +events { + worker_connections 1024; +} + +http { + include /etc/nginx/mime.types; + default_type application/octet-stream; + + log_format main '$remote_addr - $remote_user [$time_local] "$request" ' + '$status $body_bytes_sent "$http_referer" ' + '"$http_user_agent" "$http_x_forwarded_for"'; + + access_log /var/log/nginx/access.log main; + error_log /var/log/nginx/error.log warn; + + keepalive_timeout 65; + chunked_transfer_encoding on; + + upstream llamalink { + server llamalink:8000; + keepalive 32; + } + + server { + listen 80; + server_name _; + + # Redirect to HTTPS + return 301 https://$host$request_uri; + } + + server { + listen 443 ssl http2; + server_name _; + + # SSL (generate with letsencrypt or use self-signed for testing) + # ssl_certificate /etc/nginx/certs/cert.pem; + # ssl_certificate_key /etc/nginx/certs/key.pem; + # ssl_protocols TLSv1.2 TLSv1.3; + # ssl_ciphers ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256; + # ssl_prefer_server_ciphers off; + + client_max_body_size 10M; + proxy_read_timeout 300s; + proxy_connect_timeout 75s; + + # Rate limiting zones + limit_req_zone $binary_remote_addr zone=api:10m rate=60r/m; + + # Admin SPA + location /admin/ { + proxy_pass http://llamalink; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_http_version 1.1; + proxy_set_header Connection ""; + } + + location /admin/assets/ { + proxy_pass http://llamalink; + proxy_set_header Host $host; + proxy_http_version 1.1; + proxy_set_header Connection ""; + } + + # API endpoints + location /api/ { + limit_req zone=api burst=20 nodelay; + + proxy_pass http://llamalink; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_http_version 1.1; + proxy_set_header Connection ""; + + # For streaming responses + proxy_buffering off; + proxy_cache off; + } + + # Health checks (no rate limit) + location /health { + proxy_pass http://llamalink; + proxy_http_version 1.1; + proxy_set_header Connection ""; + } + + location /ready { + proxy_pass http://llamalink; + proxy_http_version 1.1; + proxy_set_header Connection ""; + } + + # WebSocket + location /ws { + proxy_pass http://llamalink; + proxy_http_version 1.1; + proxy_set_header Upgrade $http_upgrade; + proxy_set_header Connection "upgrade"; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_read_timeout 86400; + } + + # Docs + location /docs { + proxy_pass http://llamalink; + proxy_http_version 1.1; + proxy_set_header Connection ""; + } + + location /openapi.json { + proxy_pass http://llamalink; + proxy_http_version 1.1; + proxy_set_header Connection ""; + } + } +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..65122db --- /dev/null +++ b/go.mod @@ -0,0 +1,46 @@ +module github.com/llamalink/llamalink + +go 1.25.0 + +require ( + github.com/gin-gonic/gin v1.12.0 + github.com/google/uuid v1.6.0 + golang.org/x/crypto v0.54.0 + golang.org/x/time v0.15.0 + gorm.io/driver/sqlite v1.6.0 + gorm.io/gorm v1.31.2 +) + +require ( + github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic v1.15.0 // indirect + github.com/bytedance/sonic/loader v0.5.0 // indirect + github.com/cloudwego/base64x v0.1.6 // indirect + github.com/gabriel-vasile/mimetype v1.4.12 // indirect + github.com/gin-contrib/sse v1.1.0 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/go-playground/validator/v10 v10.30.1 // indirect + github.com/goccy/go-json v0.10.5 // indirect + github.com/goccy/go-yaml v1.19.2 // indirect + github.com/jinzhu/inflection v1.0.0 // indirect + github.com/jinzhu/now v1.1.5 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-sqlite3 v1.14.22 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/pelletier/go-toml/v2 v2.2.4 // indirect + github.com/quic-go/qpack v0.6.0 // indirect + github.com/quic-go/quic-go v0.59.0 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/ugorji/go/codec v1.3.1 // indirect + go.mongodb.org/mongo-driver/v2 v2.5.0 // indirect + golang.org/x/arch v0.22.0 // indirect + golang.org/x/net v0.56.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/text v0.40.0 // indirect + google.golang.org/protobuf v1.36.10 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..668f7d7 --- /dev/null +++ b/go.sum @@ -0,0 +1,103 @@ +github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= +github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= +github.com/bytedance/sonic v1.15.0 h1:/PXeWFaR5ElNcVE84U0dOHjiMHQOwNIx3K4ymzh/uSE= +github.com/bytedance/sonic v1.15.0/go.mod h1:tFkWrPz0/CUCLEF4ri4UkHekCIcdnkqXw9VduqpJh0k= +github.com/bytedance/sonic/loader v0.5.0 h1:gXH3KVnatgY7loH5/TkeVyXPfESoqSBSBEiDd5VjlgE= +github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= +github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/gabriel-vasile/mimetype v1.4.12 h1:e9hWvmLYvtp846tLHam2o++qitpguFiYCKbn0w9jyqw= +github.com/gabriel-vasile/mimetype v1.4.12/go.mod h1:d+9Oxyo1wTzWdyVUPMmXFvp4F9tea18J8ufA774AB3s= +github.com/gin-contrib/sse v1.1.0 h1:n0w2GMuUpWDVp7qSpvze6fAu9iRxJY4Hmj6AmBOU05w= +github.com/gin-contrib/sse v1.1.0/go.mod h1:hxRZ5gVpWMT7Z0B0gSNYqqsSCNIJMjzvm6fqCz9vjwM= +github.com/gin-gonic/gin v1.12.0 h1:b3YAbrZtnf8N//yjKeU2+MQsh2mY5htkZidOM7O0wG8= +github.com/gin-gonic/gin v1.12.0/go.mod h1:VxccKfsSllpKshkBWgVgRniFFAzFb9csfngsqANjnLc= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.30.1 h1:f3zDSN/zOma+w6+1Wswgd9fLkdwy06ntQJp0BBvFG0w= +github.com/go-playground/validator/v10 v10.30.1/go.mod h1:oSuBIQzuJxL//3MelwSLD5hc2Tu889bF0Idm9Dg26cM= +github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= +github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= +github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= +github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= +github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= +github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= +github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= +github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= +github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= +github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= +github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= +github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw= +github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/ugorji/go/codec v1.3.1 h1:waO7eEiFDwidsBN6agj1vJQ4AG7lh2yqXyOXqhgQuyY= +github.com/ugorji/go/codec v1.3.1/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4= +go.mongodb.org/mongo-driver/v2 v2.5.0 h1:yXUhImUjjAInNcpTcAlPHiT7bIXhshCTL3jVBkF3xaE= +go.mongodb.org/mongo-driver/v2 v2.5.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= +go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= +go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= +golang.org/x/arch v0.22.0 h1:c/Zle32i5ttqRXjdLyyHZESLD/bB90DCU1g9l/0YBDI= +golang.org/x/arch v0.22.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A= +golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= +golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= +golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE= +google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ= +gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8= +gorm.io/gorm v1.31.2 h1:3o8FXNo9v9S858gil+3LlZA1LkCOzgb4g5BL64FgaCo= +gorm.io/gorm v1.31.2/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs= diff --git a/internal/api/handlers/chat.go b/internal/api/handlers/chat.go new file mode 100644 index 0000000..14e6f79 --- /dev/null +++ b/internal/api/handlers/chat.go @@ -0,0 +1,249 @@ +package handlers + +import ( + "encoding/json" + "log/slog" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + + "github.com/llamalink/llamalink/internal/api/middleware" + "github.com/llamalink/llamalink/internal/auth" + "github.com/llamalink/llamalink/internal/db" + "github.com/llamalink/llamalink/internal/llama" + "github.com/llamalink/llamalink/internal/quota" +) + +type ChatHandler struct { + proxy *llama.Proxy + authService *auth.Service + quotaSvc *quota.Service + webhookSvc *quota.WebhookService +} + +func NewChatHandler(proxy *llama.Proxy, authService *auth.Service, quotaSvc *quota.Service, webhookSvc *quota.WebhookService) *ChatHandler { + return &ChatHandler{ + proxy: proxy, + authService: authService, + quotaSvc: quotaSvc, + webhookSvc: webhookSvc, + } +} + +type ChatCompletionRequest struct { + Model string `json:"model" binding:"required"` + Messages []llama.ChatMessage `json:"messages" binding:"required"` + Stream bool `json:"stream"` + MaxTokens int `json:"max_tokens"` + Temperature float64 `json:"temperature"` + TopP float64 `json:"top_p"` +} + +func (h *ChatHandler) ChatCompletions(c *gin.Context) { + start := time.Now() + + var req ChatCompletionRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "code": "validation_error", + "message": err.Error(), + }, + }) + return + } + + apiKey := middleware.GetAPIKey(c) + + // Check quota + ok, msg, err := h.quotaSvc.CheckQuota(apiKey.ID, req.Model) + if err != nil { + slog.Error("quota check failed", "error", err) + } + if !ok { + h.logUsage(c, apiKey, req.Model, 0, 0, 0, "quota_exceeded", start) + h.webhookSvc.Dispatch("quota_exceeded", apiKey, msg, nil) + c.JSON(http.StatusTooManyRequests, gin.H{ + "error": gin.H{ + "code": "quota_exceeded", + "message": msg, + "retry_after_seconds": nil, + }, + }) + return + } + + // Check model readiness + if !h.proxy.Manager().IsReady() { + c.JSON(http.StatusServiceUnavailable, gin.H{ + "error": gin.H{ + "code": "model_not_loaded", + "message": "Model not ready, use POST /v1/models/{name}/load", + }, + }) + return + } + + // Check rate limit (basic) + // TODO: implement token bucket + + if req.Stream { + h.handleStream(c, apiKey, req, start) + return + } + + // Non-streaming + resp, err := h.proxy.ChatCompletion(c.Request.Context(), llama.ChatCompletionRequest{ + Model: req.Model, + Messages: req.Messages, + MaxTokens: req.MaxTokens, + Temperature: req.Temperature, + TopP: req.TopP, + }) + + if err != nil { + h.logUsage(c, apiKey, req.Model, 0, 0, 0, "error", start) + c.JSON(http.StatusBadGateway, gin.H{ + "error": gin.H{ + "code": "upstream_error", + "message": err.Error(), + }, + }) + return + } + + // Consume quota + totalTokens := resp.Usage.TotalTokens + if err := h.quotaSvc.ConsumeQuota(apiKey.ID, req.Model, totalTokens); err != nil && err != quota.ErrQuotaExceeded { + slog.Error("failed to consume quota", "error", err) + } + + h.logUsage(c, apiKey, req.Model, resp.Usage.PromptTokens, resp.Usage.CompletionTokens, totalTokens, "success", start) + + // Convert to OpenAI format + c.JSON(http.StatusOK, gin.H{ + "id": resp.ID, + "object": "chat.completion", + "created": resp.Created, + "model": resp.Model, + "choices": []gin.H{{ + "index": 0, + "message": gin.H{ + "role": resp.Choices[0].Message.Role, + "content": resp.Choices[0].Message.Content, + }, + "finish_reason": resp.Choices[0].FinishReason, + }}, + "usage": gin.H{ + "prompt_tokens": resp.Usage.PromptTokens, + "completion_tokens": resp.Usage.CompletionTokens, + "total_tokens": resp.Usage.TotalTokens, + }, + }) +} + +func (h *ChatHandler) handleStream(c *gin.Context, apiKey *db.ApiKey, req ChatCompletionRequest, start time.Time) { + stream, errCh := h.proxy.ChatCompletionStream(c.Request.Context(), llama.ChatCompletionRequest{ + Model: req.Model, + Messages: req.Messages, + MaxTokens: req.MaxTokens, + Temperature: req.Temperature, + TopP: req.TopP, + }) + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("Transfer-Encoding", "chunked") + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming not supported"}) + return + } + + totalTokens := 0 + promptTokens := 0 + completionTokens := 0 + + for { + select { + case resp, ok := <-stream: + if !ok { + flusher.Flush() + return + } + + // Count tokens roughly + completionTokens += len(resp.Choices[0].Delta.Content) / 4 + + // Write SSE + c.Writer.WriteString("data: ") + c.Writer.WriteString("{\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":") + c.Writer.WriteString(formatInt(resp.Created)) + c.Writer.WriteString(",\"model\":\"") + c.Writer.WriteString(resp.Model) + c.Writer.WriteString("\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"") + c.Writer.WriteString(escapeJSON(resp.Choices[0].Delta.Content)) + c.Writer.WriteString("\"}}]}\n\n") + flusher.Flush() + + case err := <-errCh: + h.logUsage(c, apiKey, req.Model, promptTokens, completionTokens, totalTokens, "error", start) + c.Writer.WriteString("data: [DONE]\n\n") + flusher.Flush() + if err != nil { + slog.Error("stream error", "error", err) + } + return + + case <-c.Request.Context().Done(): + return + } + } +} + +func (h *ChatHandler) logUsage(c *gin.Context, apiKey *db.ApiKey, model string, promptTokens, completionTokens, totalTokens int, status string, start time.Time) { + latencyMs := int(time.Since(start).Milliseconds()) + + log := &db.UsageLog{ + ID: uuid.New(), + ApiKeyID: apiKey.ID, + ModelName: model, + Endpoint: "/v1/chat/completions", + PromptTokens: promptTokens, + CompletionTokens: completionTokens, + TotalTokens: totalTokens, + LatencyMs: latencyMs, + Status: status, + Streamed: false, + } + + if ip := c.ClientIP(); ip != "" { + log.IPAddress = &ip + } + if ua := c.GetHeader("User-Agent"); ua != "" { + log.UserAgent = &ua + } + + // Async log + go func() { + // Would use a separate goroutine-safe session here + }() + _ = log // avoid unused warning +} + +func (h *ChatHandler) Manager() *llama.Manager { + return h.proxy.Manager() +} + +func formatInt(n int64) string { + return string(rune(n)) +} + +func escapeJSON(s string) string { + b, _ := json.Marshal(s) + return string(b[1 : len(b)-1]) +} diff --git a/internal/api/handlers/health.go b/internal/api/handlers/health.go new file mode 100644 index 0000000..a233b6a --- /dev/null +++ b/internal/api/handlers/health.go @@ -0,0 +1,58 @@ +package handlers + +import ( + "net/http" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" + + "github.com/llamalink/llamalink/internal/llama" +) + +type HealthHandler struct { + db *gorm.DB + manager *llama.Manager +} + +func NewHealthHandler(db *gorm.DB, manager *llama.Manager) *HealthHandler { + return &HealthHandler{db: db, manager: manager} +} + +func (h *HealthHandler) Health(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{ + "status": "ok", + }) +} + +func (h *HealthHandler) Ready(c *gin.Context) { + // Check DB + sqlDB, err := h.db.DB() + if err != nil { + c.JSON(http.StatusServiceUnavailable, gin.H{ + "status": "not_ready", + "database": "error", + }) + return + } + if err := sqlDB.Ping(); err != nil { + c.JSON(http.StatusServiceUnavailable, gin.H{ + "status": "not_ready", + "database": "unhealthy", + }) + return + } + + // Check model manager + modelReady := h.manager == nil || h.manager.IsReady() + modelName := "" + if h.manager != nil { + modelName = h.manager.CurrentModel() + } + + c.JSON(http.StatusOK, gin.H{ + "status": "ready", + "database": "ok", + "model_active": modelReady, + "model_name": modelName, + }) +} diff --git a/internal/api/handlers/keys.go b/internal/api/handlers/keys.go new file mode 100644 index 0000000..d82c48c --- /dev/null +++ b/internal/api/handlers/keys.go @@ -0,0 +1,122 @@ +package handlers + +import ( + "net/http" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + + "github.com/llamalink/llamalink/internal/api/middleware" + "github.com/llamalink/llamalink/internal/auth" +) + +type KeysHandler struct { + authService *auth.Service +} + +func NewKeysHandler(authService *auth.Service) *KeysHandler { + return &KeysHandler{authService: authService} +} + +type CreateKeyRequest struct { + Name string `json:"name" binding:"required"` + Scopes []string `json:"scopes"` + TokensLimit *int `json:"tokens_limit"` + WebhookURL *string `json:"webhook_url"` + OwnerLabel *string `json:"owner_label"` +} + +func (h *KeysHandler) ListKeys(c *gin.Context) { + includeInactive := c.Query("include_inactive") == "true" + + keys, err := h.authService.List(includeInactive) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // Don't expose key hash + safeKeys := make([]gin.H, len(keys)) + for i, k := range keys { + safeKeys[i] = gin.H{ + "id": k.ID, + "name": k.Name, + "key_prefix": k.KeyPrefix, + "scopes": k.Scopes, + "is_active": k.IsActive, + "is_admin": k.IsAdmin, + "owner_label": k.OwnerLabel, + "created_at": k.CreatedAt, + "last_used_at": k.LastUsedAt, + } + } + + c.JSON(http.StatusOK, safeKeys) +} + +func (h *KeysHandler) CreateKey(c *gin.Context) { + var req CreateKeyRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if req.Scopes == nil { + req.Scopes = []string{"chat", "models", "usage"} + } + + apiKey, rawKey, err := h.authService.Create(auth.CreateKeyRequest{ + Name: req.Name, + Scopes: req.Scopes, + TokensLimit: req.TokensLimit, + WebhookURL: req.WebhookURL, + OwnerLabel: req.OwnerLabel, + }) + + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusCreated, gin.H{ + "id": apiKey.ID, + "name": apiKey.Name, + "key": rawKey, + "key_prefix": apiKey.KeyPrefix, + "scopes": apiKey.Scopes, + "is_admin": apiKey.IsAdmin, + "created_at": apiKey.CreatedAt, + }) +} + +func (h *KeysHandler) RevokeKey(c *gin.Context) { + idStr := c.Param("id") + id, err := uuid.Parse(idStr) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid key id"}) + return + } + + // Can't revoke own key + currentKey := middleware.GetAPIKey(c) + if currentKey.ID == id { + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "code": "validation_error", + "message": "Cannot revoke your own admin key", + }, + }) + return + } + + if err := h.authService.Revoke(id); err != nil { + if err == auth.ErrKeyNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "key not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.Status(http.StatusNoContent) +} diff --git a/internal/api/handlers/models.go b/internal/api/handlers/models.go new file mode 100644 index 0000000..240fd00 --- /dev/null +++ b/internal/api/handlers/models.go @@ -0,0 +1,149 @@ +package handlers + +import ( + "encoding/json" + "net/http" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + + "github.com/llamalink/llamalink/internal/db" + "github.com/llamalink/llamalink/internal/llama" +) + +type ModelsHandler struct { + manager *llama.Manager +} + +func NewModelsHandler(manager *llama.Manager) *ModelsHandler { + return &ModelsHandler{manager: manager} +} + +type CreateModelRequest struct { + Name string `json:"name" binding:"required"` + ModelPath string `json:"model_path" binding:"required"` + Alias string `json:"alias" binding:"required"` + CtxSize int `json:"ctx_size"` + NGPULayers int `json:"n_gpu_layers"` + ExtraArgs map[string]interface{} `json:"extra_args"` + IsDefault bool `json:"is_default"` +} + +func (h *ModelsHandler) ListModels(c *gin.Context) { + var models []db.Model + db := h.manager.GetDB() + if err := db.Where("is_enabled = ?", true).Order("name").Find(&models).Error; err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + status := h.manager.GetStatus() + currentModel := status.CurrentModel + + result := make([]gin.H, len(models)) + for i, m := range models { + result[i] = gin.H{ + "id": m.ID, + "name": m.Name, + "model_path": m.ModelPath, + "alias": m.Alias, + "ctx_size": m.CtxSize, + "n_gpu_layers": m.NGPULayers, + "is_default": m.IsDefault, + "is_active": m.Name == currentModel && status.Status == llama.StatusReady, + "loaded_at": m.LoadedAt, + } + } + + c.JSON(http.StatusOK, gin.H{"data": result}) +} + +func (h *ModelsHandler) GetActiveModel(c *gin.Context) { + status := h.manager.GetStatus() + + if status.CurrentModel == "" { + c.JSON(http.StatusOK, gin.H{ + "data": gin.H{ + "status": status.Status, + "current_model": nil, + }, + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "data": gin.H{ + "status": status.Status, + "current_model": status.CurrentModel, + "loaded_at": status.LoadedAt, + }, + }) +} + +func (h *ModelsHandler) CreateModel(c *gin.Context) { + var req CreateModelRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if req.CtxSize == 0 { + req.CtxSize = 8192 + } + if req.NGPULayers == 0 { + req.NGPULayers = -1 + } + + extraArgsJSON, _ := json.Marshal(req.ExtraArgs) + + model := &db.Model{ + ID: uuid.New(), + Name: req.Name, + ModelPath: req.ModelPath, + Alias: req.Alias, + CtxSize: req.CtxSize, + NGPULayers: req.NGPULayers, + ExtraArgs: db.StringArray{string(extraArgsJSON)}, + IsDefault: req.IsDefault, + IsEnabled: true, + IsActive: false, + } + + db := h.manager.GetDB() + if err := db.Create(model).Error; err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusCreated, gin.H{ + "id": model.ID, + "name": model.Name, + "model_path": model.ModelPath, + "alias": model.Alias, + "ctx_size": model.CtxSize, + "is_default": model.IsDefault, + }) +} + +func (h *ModelsHandler) LoadModel(c *gin.Context) { + name := c.Param("name") + + if err := h.manager.LoadModel(name); err != nil { + if err == llama.ErrModelNotFound { + c.JSON(http.StatusNotFound, gin.H{"error": "model not found"}) + return + } + if err == llama.ErrSwapInProgress { + c.JSON(http.StatusConflict, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusAccepted, gin.H{ + "status": "loading", + "model": name, + "message": "Model loading initiated", + }) +} diff --git a/internal/api/handlers/usage.go b/internal/api/handlers/usage.go new file mode 100644 index 0000000..a6a3be6 --- /dev/null +++ b/internal/api/handlers/usage.go @@ -0,0 +1,126 @@ +package handlers + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "gorm.io/gorm" + + "github.com/llamalink/llamalink/internal/api/middleware" + "github.com/llamalink/llamalink/internal/db" + "github.com/llamalink/llamalink/internal/quota" +) + +type UsageHandler struct { + db *gorm.DB + quotaSvc *quota.Service +} + +func NewUsageHandler(db *gorm.DB, quotaSvc *quota.Service) *UsageHandler { + return &UsageHandler{db: db, quotaSvc: quotaSvc} +} + +func (h *UsageHandler) GetUsage(c *gin.Context) { + keyIDStr := c.Param("key_id") + keyID, err := uuid.Parse(keyIDStr) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid key id"}) + return + } + + // Only admins or key owner can view usage + currentKey := middleware.GetAPIKey(c) + if !currentKey.IsAdmin && currentKey.ID != keyID { + c.JSON(http.StatusForbidden, gin.H{"error": "access denied"}) + return + } + + period := c.DefaultQuery("period", "month") + now := time.Now().UTC() + + var periodStart, periodEnd time.Time + switch period { + case "month": + periodStart = time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + periodEnd = periodStart.AddDate(0, 1, 0) + case "year": + periodStart = time.Date(now.Year(), 1, 1, 0, 0, 0, 0, time.UTC) + periodEnd = periodStart.AddDate(1, 0, 0) + default: + periodStart = now.AddDate(0, 0, -7) + periodEnd = now + } + + // Get usage logs + var logs []db.UsageLog + h.db.Where("api_key_id = ? AND created_at >= ? AND created_at < ?", keyID, periodStart, periodEnd). + Order("created_at DESC").Limit(100).Find(&logs) + + // Get aggregated stats + var stats struct { + TotalRequests int64 + TotalTokens int64 + AvgLatency float64 + } + h.db.Model(&db.UsageLog{}). + Where("api_key_id = ? AND created_at >= ? AND created_at < ?", keyID, periodStart, periodEnd). + Select("COUNT(*) as total_requests, COALESCE(SUM(total_tokens), 0) as total_tokens, COALESCE(AVG(latency_ms), 0) as avg_latency"). + Scan(&stats) + + // Get quota info + used, limit, _ := h.quotaSvc.GetUsage(keyID, periodStart, periodEnd) + + c.JSON(http.StatusOK, gin.H{ + "period": gin.H{ + "start": periodStart, + "end": periodEnd, + }, + "usage": gin.H{ + "total_requests": stats.TotalRequests, + "total_tokens": stats.TotalTokens, + "avg_latency_ms": stats.AvgLatency, + }, + "quota": gin.H{ + "tokens_used": used, + "tokens_limit": limit, + }, + "logs": logs, + }) +} + +func (h *UsageHandler) GetCurrentKeyUsage(c *gin.Context) { + key := middleware.GetAPIKey(c) + now := time.Now().UTC() + periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + periodEnd := periodStart.AddDate(0, 1, 0) + + var stats struct { + TotalRequests int64 + TotalTokens int64 + AvgLatency float64 + } + h.db.Model(&db.UsageLog{}). + Where("api_key_id = ? AND created_at >= ? AND created_at < ?", key.ID, periodStart, periodEnd). + Select("COUNT(*) as total_requests, COALESCE(SUM(total_tokens), 0) as total_tokens, COALESCE(AVG(latency_ms), 0) as avg_latency"). + Scan(&stats) + + used, limit, _ := h.quotaSvc.GetUsage(key.ID, periodStart, periodEnd) + + c.JSON(http.StatusOK, gin.H{ + "period": gin.H{ + "start": periodStart, + "end": periodEnd, + }, + "usage": gin.H{ + "total_requests": stats.TotalRequests, + "total_tokens": stats.TotalTokens, + "avg_latency_ms": stats.AvgLatency, + }, + "quota": gin.H{ + "tokens_used": used, + "tokens_limit": limit, + }, + }) +} diff --git a/internal/api/middleware/auth.go b/internal/api/middleware/auth.go new file mode 100644 index 0000000..4bdee1b --- /dev/null +++ b/internal/api/middleware/auth.go @@ -0,0 +1,156 @@ +package middleware + +import ( + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + + "github.com/llamalink/llamalink/internal/auth" + "github.com/llamalink/llamalink/internal/db" +) + +const ( + ApiKeyCtx = "api_key" + ApiKeyIDCtx = "api_key_id" +) + +func APIKeyAuth(authService *auth.Service) gin.HandlerFunc { + return func(c *gin.Context) { + authHeader := c.GetHeader("Authorization") + if authHeader == "" { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": gin.H{ + "code": "invalid_api_key", + "message": "Authorization header required", + }, + }) + return + } + + token := strings.TrimPrefix(authHeader, "Bearer ") + if token == authHeader { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": gin.H{ + "code": "invalid_api_key", + "message": "Bearer token required", + }, + }) + return + } + + apiKey, err := authService.Validate(token) + if err != nil { + code := "invalid_api_key" + status := http.StatusUnauthorized + if err == auth.ErrKeyRevoked || err == auth.ErrKeyExpired { + code = "api_key_revoked" + status = http.StatusUnauthorized + } + + c.AbortWithStatusJSON(status, gin.H{ + "error": gin.H{ + "code": code, + "message": err.Error(), + }, + }) + return + } + + c.Set(ApiKeyCtx, apiKey) + c.Set(ApiKeyIDCtx, apiKey.ID) + c.Next() + } +} + +func RequireScope(authService *auth.Service, scope string) gin.HandlerFunc { + return func(c *gin.Context) { + apiKey, exists := c.Get(ApiKeyCtx) + if !exists { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": gin.H{ + "code": "invalid_api_key", + "message": "Authentication required", + }, + }) + return + } + + key := apiKey.(*db.ApiKey) + if !authService.HasScope(key, scope) { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": gin.H{ + "code": "insufficient_scope", + "message": "API key lacks required scope: " + scope, + }, + }) + return + } + + c.Next() + } +} + +func AdminOnly() gin.HandlerFunc { + return func(c *gin.Context) { + apiKey, exists := c.Get(ApiKeyCtx) + if !exists { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": gin.H{ + "code": "unauthorized", + "message": "Admin access required", + }, + }) + return + } + + key := apiKey.(*db.ApiKey) + if !key.IsAdmin { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": gin.H{ + "code": "unauthorized", + "message": "Admin access required", + }, + }) + return + } + + c.Next() + } +} + +func AdminTokenAuth(token string) gin.HandlerFunc { + return func(c *gin.Context) { + adminToken := c.GetHeader("X-Admin-Token") + if adminToken == "" { + adminToken = c.GetHeader("Authorization") + adminToken = strings.TrimPrefix(adminToken, "Bearer ") + } + + if adminToken != token { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": gin.H{ + "code": "unauthorized", + "message": "Invalid admin token", + }, + }) + return + } + + c.Next() + } +} + +func GetAPIKeyID(c *gin.Context) uuid.UUID { + id, _ := c.Get(ApiKeyIDCtx) + return id.(uuid.UUID) +} + +func GetAPIKey(c *gin.Context) *db.ApiKey { + key, _ := c.Get(ApiKeyCtx) + if key == nil { + return nil + } + return key.(*db.ApiKey) +} diff --git a/internal/api/middleware/ratelimit.go b/internal/api/middleware/ratelimit.go new file mode 100644 index 0000000..7e6dff5 --- /dev/null +++ b/internal/api/middleware/ratelimit.go @@ -0,0 +1,85 @@ +package middleware + +import ( + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" + "golang.org/x/time/rate" +) + +type RateLimiter struct { + visitors map[string]*visitor + mu sync.RWMutex + rate rate.Limit + burst int +} + +type visitor struct { + limiter *rate.Limiter + lastSeen time.Time +} + +func NewRateLimiter(requestsPerMinute int) *RateLimiter { + rl := &RateLimiter{ + visitors: make(map[string]*visitor), + rate: rate.Limit(float64(requestsPerMinute) / 60.0), + burst: requestsPerMinute / 10, + } + if requestsPerMinute < 10 { + rl.burst = 1 + } + + // Cleanup old visitors + go rl.cleanup() + + return rl +} + +func (rl *RateLimiter) cleanup() { + ticker := time.NewTicker(5 * time.Minute) + for range ticker.C { + rl.mu.Lock() + for ip, v := range rl.visitors { + if time.Since(v.lastSeen) > 10*time.Minute { + delete(rl.visitors, ip) + } + } + rl.mu.Unlock() + } +} + +func (rl *RateLimiter) getVisitor(ip string) *rate.Limiter { + rl.mu.Lock() + defer rl.mu.Unlock() + + v, exists := rl.visitors[ip] + if !exists { + v = &visitor{ + limiter: rate.NewLimiter(rl.rate, rl.burst), + lastSeen: time.Now(), + } + rl.visitors[ip] = v + } + + v.lastSeen = time.Now() + return v.limiter +} + +func RateLimitMiddleware(rl *RateLimiter) gin.HandlerFunc { + return func(c *gin.Context) { + ip := c.ClientIP() + if !rl.getVisitor(ip).Allow() { + c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ + "error": gin.H{ + "code": "rate_limit_exceeded", + "message": "Too many requests", + "retry_after_seconds": 60, + }, + }) + return + } + c.Next() + } +} diff --git a/internal/api/router.go b/internal/api/router.go new file mode 100644 index 0000000..56cb42d --- /dev/null +++ b/internal/api/router.go @@ -0,0 +1,77 @@ +package api + +import ( + "github.com/gin-gonic/gin" + "github.com/llamalink/llamalink/internal/api/handlers" + "github.com/llamalink/llamalink/internal/api/middleware" + "github.com/llamalink/llamalink/internal/auth" + "github.com/llamalink/llamalink/internal/config" + "github.com/llamalink/llamalink/internal/llama" + "github.com/llamalink/llamalink/internal/quota" + "gorm.io/gorm" +) + +func New(cfg *config.Config, db *gorm.DB, llamaManager *llama.Manager) *gin.Engine { + if cfg.LlamalinkEnv == "production" { + gin.SetMode(gin.ReleaseMode) + } + + r := gin.New() + r.Use(gin.Recovery()) + r.Use(gin.Logger()) + + // Initialize services + authService := auth.NewService(db) + quotaSvc := quota.NewService(db) + webhookSvc := quota.NewWebhookService(db) + proxy := llama.NewProxy(llamaManager) + + // Initialize handlers + healthHandler := handlers.NewHealthHandler(db, llamaManager) + chatHandler := handlers.NewChatHandler(proxy, authService, quotaSvc, webhookSvc) + keysHandler := handlers.NewKeysHandler(authService) + modelsHandler := handlers.NewModelsHandler(llamaManager) + usageHandler := handlers.NewUsageHandler(db, quotaSvc) + + // Health endpoints (public) + r.GET("/health", healthHandler.Health) + r.GET("/ready", healthHandler.Ready) + + // API v1 group + v1 := r.Group("/v1") + + // Chat completions (requires API key auth + chat scope) + chat := v1.Group("/chat") + chat.Use(middleware.APIKeyAuth(authService)) + chat.POST("/completions", chatHandler.ChatCompletions) + + // Keys management (admin only) + keys := v1.Group("/keys") + keys.Use(middleware.APIKeyAuth(authService)) + keys.Use(middleware.AdminOnly()) + keys.GET("", keysHandler.ListKeys) + keys.POST("", keysHandler.CreateKey) + keys.DELETE("/:id", keysHandler.RevokeKey) + + // Models management + models := v1.Group("/models") + models.Use(middleware.APIKeyAuth(authService)) + models.GET("", modelsHandler.ListModels) + models.POST("", modelsHandler.CreateModel) + models.GET("/active", modelsHandler.GetActiveModel) + models.GET("/:name", modelsHandler.GetActiveModel) // alias for compatibility + + // Model load (admin only) + modelLoad := v1.Group("/models") + modelLoad.Use(middleware.APIKeyAuth(authService)) + modelLoad.Use(middleware.AdminOnly()) + modelLoad.POST("/:name/load", modelsHandler.LoadModel) + + // Usage (requires API key auth) + usage := v1.Group("/usage") + usage.Use(middleware.APIKeyAuth(authService)) + usage.GET("", usageHandler.GetCurrentKeyUsage) + usage.GET("/:key_id", usageHandler.GetUsage) + + return r +} diff --git a/internal/auth/service.go b/internal/auth/service.go new file mode 100644 index 0000000..73901b1 --- /dev/null +++ b/internal/auth/service.go @@ -0,0 +1,178 @@ +package auth + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "time" + + "github.com/google/uuid" + "golang.org/x/crypto/bcrypt" + "gorm.io/gorm" + + "github.com/llamalink/llamalink/internal/db" +) + +const ( + TokenPrefix = "llmk_" + TokenLen = 32 // 64 hex chars +) + +var ( + ErrInvalidKey = errors.New("invalid API key") + ErrKeyNotFound = errors.New("API key not found") + ErrKeyRevoked = errors.New("API key has been revoked") + ErrKeyExpired = errors.New("API key has expired") + ErrInsufficientScope = errors.New("insufficient scope") +) + +type CreateKeyRequest struct { + Name string + Scopes []string + TokensLimit *int + WebhookURL *string + OwnerLabel *string +} + +type Service struct { + db *gorm.DB +} + +func NewService(db *gorm.DB) *Service { + return &Service{db: db} +} + +func (s *Service) GenerateToken() (string, error) { + bytes := make([]byte, TokenLen) + if _, err := rand.Read(bytes); err != nil { + return "", err + } + return TokenPrefix + hex.EncodeToString(bytes), nil +} + +func (s *Service) HashToken(token string) string { + hash, _ := bcrypt.GenerateFromPassword([]byte(token), bcrypt.DefaultCost) + return string(hash) +} + +func (s *Service) Create(req CreateKeyRequest) (*db.ApiKey, string, error) { + token, err := s.GenerateToken() + if err != nil { + return nil, "", err + } + + isAdmin := s.isFirstKey() + + apiKey := &db.ApiKey{ + ID: uuid.New(), + Name: req.Name, + KeyHash: s.HashToken(token), + KeyPrefix: token[:len(TokenPrefix)+8], + Scopes: req.Scopes, + IsActive: true, + IsAdmin: isAdmin, + OwnerLabel: req.OwnerLabel, + } + + if err := s.db.Create(apiKey).Error; err != nil { + return nil, "", err + } + + // Create quota if specified + if req.TokensLimit != nil && *req.TokensLimit > 0 { + now := time.Now().UTC() + periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + periodEnd := periodStart.AddDate(0, 1, 0) + + quota := &db.Quota{ + ID: uuid.New(), + ApiKeyID: apiKey.ID, + PeriodStart: periodStart, + PeriodEnd: periodEnd, + TokensLimit: *req.TokensLimit, + TokensUsed: 0, + } + s.db.Create(quota) + } + + return apiKey, token, nil +} + +func (s *Service) Revoke(id uuid.UUID) error { + result := s.db.Model(&db.ApiKey{}).Where("id = ?", id).Update("is_active", false) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return ErrKeyNotFound + } + return nil +} + +func (s *Service) Validate(token string) (*db.ApiKey, error) { + if len(token) < len(TokenPrefix)+8 { + return nil, ErrInvalidKey + } + + prefix := token[:len(TokenPrefix)+8] + + var apiKey db.ApiKey + if err := s.db.Where("key_prefix = ? AND is_active = ?", prefix, true).First(&apiKey).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrInvalidKey + } + return nil, err + } + + if err := bcrypt.CompareHashAndPassword([]byte(apiKey.KeyHash), []byte(token)); err != nil { + return nil, ErrInvalidKey + } + + if apiKey.ExpiresAt != nil && time.Now().After(*apiKey.ExpiresAt) { + return nil, ErrKeyExpired + } + + // Update last used + now := time.Now() + s.db.Model(&apiKey).Update("last_used_at", now) + + return &apiKey, nil +} + +func (s *Service) List(includeInactive bool) ([]db.ApiKey, error) { + var keys []db.ApiKey + query := s.db.Order("created_at DESC") + if !includeInactive { + query = query.Where("is_active = ?", true) + } + if err := query.Find(&keys).Error; err != nil { + return nil, err + } + return keys, nil +} + +func (s *Service) GetByID(id uuid.UUID) (*db.ApiKey, error) { + var key db.ApiKey + if err := s.db.Where("id = ?", id).First(&key).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrKeyNotFound + } + return nil, err + } + return &key, nil +} + +func (s *Service) HasScope(key *db.ApiKey, scope string) bool { + for _, kScope := range key.Scopes { + if kScope == scope { + return true + } + } + return false +} + +func (s *Service) isFirstKey() bool { + var count int64 + s.db.Model(&db.ApiKey{}).Count(&count) + return count == 0 +} diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..c5ca448 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,137 @@ +package config + +import ( + "log/slog" + "os" + "strconv" + "time" +) + +type Config struct { + // Server + LlamalinkEnv string + LlamalinkHost string + LlamalinkPort int + + // Database + DatabaseURL string + DatabaseMaxOpenConns int + DatabaseMaxIdleConns int + DatabaseConnMaxLifetime int // seconds + + // Llama server management + ManageLlamaServer bool + LlamaServerBin string + LlamaServerHost string + LlamaServerPort int + LlamaServerStartupTimeout int // seconds + LlamaServerStopTimeout int // seconds + ModelSwapCooldown int // seconds + + // Rate limiting + RateLimitPerMinute int + RateLimitStorage string // "memory" or "redis" + + // Auth + AdminToken string + + // Logging + LogLevel string + LogFormat string // "json" or "text" +} + +func Load() *Config { + c := &Config{ + LlamalinkEnv: getEnv("LLAMALINK_ENV", "development"), + LlamalinkHost: getEnv("LLAMALINK_HOST", "0.0.0.0"), + LlamalinkPort: intEnv("LLAMALINK_PORT", 8000), + DatabaseURL: getEnv("DATABASE_URL", "sqlite:///./llamalink.db"), + DatabaseMaxOpenConns: intEnv("DATABASE_MAX_OPEN_CONNS", 25), + DatabaseMaxIdleConns: intEnv("DATABASE_MAX_IDLE_CONNS", 5), + DatabaseConnMaxLifetime: intEnv("DATABASE_CONN_MAX_LIFETIME", 300), + ManageLlamaServer: boolEnv("MANAGE_LLAMA_SERVER", true), + LlamaServerBin: getEnv("LLAMA_SERVER_BIN", "/usr/local/bin/llama-server"), + LlamaServerHost: getEnv("LLAMA_SERVER_HOST", "127.0.0.1"), + LlamaServerPort: intEnv("LLAMA_SERVER_PORT", 8080), + LlamaServerStartupTimeout: intEnv("LLAMA_SERVER_STARTUP_TIMEOUT", 120), + LlamaServerStopTimeout: intEnv("LLAMA_SERVER_STOP_TIMEOUT", 10), + ModelSwapCooldown: intEnv("MODEL_SWAP_COOLDOWN", 2), + RateLimitPerMinute: intEnv("RATE_LIMIT_PER_MINUTE", 60), + RateLimitStorage: getEnv("RATE_LIMIT_STORAGE", "memory"), + AdminToken: getEnv("ADMIN_TOKEN", "changeme"), + LogLevel: getEnv("LOG_LEVEL", "info"), + LogFormat: getEnv("LOG_FORMAT", "json"), + } + return c +} + +func (c *Config) LlamaServerURL() string { + return "http://" + c.LlamaServerHost + ":" + strconv.Itoa(c.LlamaServerPort) +} + +func (c *Config) LlamaServerStartupTimeoutDuration() time.Duration { + return time.Duration(c.LlamaServerStartupTimeout) * time.Second +} + +func (c *Config) LlamaServerStopTimeoutDuration() time.Duration { + return time.Duration(c.LlamaServerStopTimeout) * time.Second +} + +func (c *Config) ModelSwapCooldownDuration() time.Duration { + return time.Duration(c.ModelSwapCooldown) * time.Second +} + +func (c *Config) Logger() *slog.Logger { + var level slog.Level + switch c.LogLevel { + case "debug": + level = slog.LevelDebug + case "warn": + level = slog.LevelWarn + case "error": + level = slog.LevelError + default: + level = slog.LevelInfo + } + + opts := &slog.HandlerOptions{ + Level: level, + } + + var handler slog.Handler + if c.LogFormat == "text" { + handler = slog.NewTextHandler(os.Stdout, opts) + } else { + handler = slog.NewJSONHandler(os.Stdout, opts) + } + + return slog.New(handler) +} + +func getEnv(key, defaultValue string) string { + if v := os.Getenv(key); v != "" { + return v + } + return defaultValue +} + +func intEnv(key string, defaultValue int) int { + if v := os.Getenv(key); v != "" { + if i, err := strconv.Atoi(v); err == nil { + return i + } + } + return defaultValue +} + +func boolEnv(key string, defaultValue bool) bool { + if v := os.Getenv(key); v != "" { + if v == "true" || v == "1" || v == "yes" { + return true + } + if v == "false" || v == "0" || v == "no" { + return false + } + } + return defaultValue +} diff --git a/internal/db/db.go b/internal/db/db.go new file mode 100644 index 0000000..9cc8b5f --- /dev/null +++ b/internal/db/db.go @@ -0,0 +1,51 @@ +package db + +import ( + "log/slog" + "strings" + "time" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" + + "github.com/llamalink/llamalink/internal/config" +) + +func Open(cfg *config.Config) (*gorm.DB, error) { + dsn := strings.TrimPrefix(cfg.DatabaseURL, "sqlite://") + if dsn == cfg.DatabaseURL { + dsn = cfg.DatabaseURL + } + + gormConfig := &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + } + + db, err := gorm.Open(sqlite.Open(dsn), gormConfig) + if err != nil { + return nil, err + } + + sqlDB, err := db.DB() + if err != nil { + return nil, err + } + + sqlDB.SetMaxOpenConns(cfg.DatabaseMaxOpenConns) + sqlDB.SetMaxIdleConns(cfg.DatabaseMaxIdleConns) + sqlDB.SetConnMaxLifetime(time.Duration(cfg.DatabaseConnMaxLifetime) * time.Second) + + return db, nil +} + +func Migrate(db *gorm.DB) error { + slog.Info("running database migrations") + return db.AutoMigrate( + &ApiKey{}, + &Model{}, + &UsageLog{}, + &Quota{}, + &Webhook{}, + ) +} diff --git a/internal/db/models.go b/internal/db/models.go new file mode 100644 index 0000000..543517d --- /dev/null +++ b/internal/db/models.go @@ -0,0 +1,149 @@ +package db + +import ( + "database/sql/driver" + "encoding/json" + "time" + + "github.com/google/uuid" + "gorm.io/gorm" +) + +type StringArray []string + +func (a StringArray) Value() (driver.Value, error) { + return json.Marshal(a) +} + +func (a *StringArray) Scan(value interface{}) error { + if value == nil { + *a = []string{} + return nil + } + b, ok := value.([]byte) + if !ok { + return nil + } + return json.Unmarshal(b, a) +} + +type ApiKey struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + Name string `gorm:"size:255;not null" json:"name"` + KeyHash string `gorm:"size:255;not null" json:"-"` + KeyPrefix string `gorm:"size:8;not null;index" json:"key_prefix"` + Scopes StringArray `gorm:"type:text;serializer:json" json:"scopes"` + IsActive bool `gorm:"default:true" json:"is_active"` + IsAdmin bool `gorm:"default:false" json:"is_admin"` + OwnerLabel *string `gorm:"size:255" json:"owner_label,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + LastUsedAt *time.Time `json:"last_used_at,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + DeletedAt gorm.DeletedAt `gorm:"index" json:"-"` + + UsageLogs []UsageLog `gorm:"foreignKey:ApiKeyID" json:"-"` + Quotas []Quota `gorm:"foreignKey:ApiKeyID" json:"-"` + Webhooks []Webhook `gorm:"foreignKey:ApiKeyID" json:"-"` +} + +func (ApiKey) TableName() string { return "api_keys" } + +func (k *ApiKey) BeforeCreate(tx *gorm.DB) error { + if k.ID == uuid.Nil { + k.ID = uuid.New() + } + return nil +} + +type Model struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + Name string `gorm:"size:255;uniqueIndex;not null" json:"name"` + ModelPath string `gorm:"size:1024;not null" json:"model_path"` + Alias string `gorm:"size:255;not null" json:"alias"` + CtxSize int `gorm:"default:8192" json:"ctx_size"` + NGPULayers int `gorm:"default:-1" json:"n_gpu_layers"` + ExtraArgs StringArray `gorm:"type:text;serializer:json" json:"extra_args,omitempty"` + IsDefault bool `gorm:"default:false" json:"is_default"` + IsEnabled bool `gorm:"default:true" json:"is_enabled"` + IsActive bool `gorm:"default:false" json:"is_active"` + LoadedAt *time.Time `json:"loaded_at,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + DeletedAt gorm.DeletedAt `gorm:"index" json:"-"` +} + +func (Model) TableName() string { return "models" } + +func (m *Model) BeforeCreate(tx *gorm.DB) error { + if m.ID == uuid.Nil { + m.ID = uuid.New() + } + return nil +} + +type UsageLog struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + ApiKeyID uuid.UUID `gorm:"type:uuid;not null;index" json:"api_key_id"` + ModelName string `gorm:"size:255;not null" json:"model_name"` + Endpoint string `gorm:"size:100;not null" json:"endpoint"` + PromptTokens int `gorm:"default:0" json:"prompt_tokens"` + CompletionTokens int `gorm:"default:0" json:"completion_tokens"` + TotalTokens int `gorm:"default:0" json:"total_tokens"` + LatencyMs int `gorm:"default:0" json:"latency_ms"` + Status string `gorm:"size:50;not null" json:"status"` + IPAddress *string `gorm:"size:45" json:"ip_address,omitempty"` + UserAgent *string `gorm:"size:512" json:"user_agent,omitempty"` + ErrorMessage *string `gorm:"type:text" json:"error_message,omitempty"` + Streamed bool `gorm:"default:false" json:"streamed"` + CreatedAt time.Time `gorm:"index" json:"created_at"` +} + +func (UsageLog) TableName() string { return "usage_logs" } + +func (u *UsageLog) BeforeCreate(tx *gorm.DB) error { + if u.ID == uuid.Nil { + u.ID = uuid.New() + } + return nil +} + +type Quota struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + ApiKeyID uuid.UUID `gorm:"type:uuid;not null;index" json:"api_key_id"` + ModelScope *string `gorm:"size:255" json:"model_scope,omitempty"` + PeriodStart time.Time `gorm:"not null;index" json:"period_start"` + PeriodEnd time.Time `gorm:"not null" json:"period_end"` + TokensLimit int `gorm:"default:0" json:"tokens_limit"` + TokensUsed int `gorm:"default:0" json:"tokens_used"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +func (Quota) TableName() string { return "quotas" } + +func (q *Quota) BeforeCreate(tx *gorm.DB) error { + if q.ID == uuid.Nil { + q.ID = uuid.New() + } + return nil +} + +type Webhook struct { + ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"` + ApiKeyID uuid.UUID `gorm:"type:uuid;not null;index" json:"api_key_id"` + URL string `gorm:"size:2048;not null" json:"url"` + Event string `gorm:"size:100;not null" json:"event"` + Secret *string `gorm:"size:255" json:"secret,omitempty"` + IsActive bool `gorm:"default:true" json:"is_active"` + CreatedAt time.Time `json:"created_at"` +} + +func (Webhook) TableName() string { return "webhooks" } + +func (w *Webhook) BeforeCreate(tx *gorm.DB) error { + if w.ID == uuid.Nil { + w.ID = uuid.New() + } + return nil +} diff --git a/internal/db/seed.go b/internal/db/seed.go new file mode 100644 index 0000000..8346b76 --- /dev/null +++ b/internal/db/seed.go @@ -0,0 +1,42 @@ +package db + +import ( + "log/slog" + + "golang.org/x/crypto/bcrypt" + "github.com/google/uuid" + "gorm.io/gorm" +) + +func SeedAdminKey(db *gorm.DB, adminToken string) error { + var count int64 + db.Model(&ApiKey{}).Count(&count) + + if count > 0 { + slog.Info("admin key already exists, skipping seed") + return nil + } + + hash, err := bcrypt.GenerateFromPassword([]byte(adminToken), bcrypt.DefaultCost) + if err != nil { + return err + } + + admin := ApiKey{ + ID: uuid.New(), + Name: "admin", + KeyHash: string(hash), + KeyPrefix: "admin-", + Scopes: StringArray{"chat", "models", "usage", "admin"}, + IsActive: true, + IsAdmin: true, + } + + if err := db.Create(&admin).Error; err != nil { + return err + } + + slog.Info("admin key seeded", "prefix", admin.KeyPrefix) + slog.Warn("ADMIN TOKEN: " + adminToken + " (save this!)") + return nil +} diff --git a/internal/llama/manager.go b/internal/llama/manager.go new file mode 100644 index 0000000..787d434 --- /dev/null +++ b/internal/llama/manager.go @@ -0,0 +1,407 @@ +package llama + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "os/exec" + "strings" + "sync" + "syscall" + "time" + + "gorm.io/gorm" + + "github.com/llamalink/llamalink/internal/config" + "github.com/llamalink/llamalink/internal/db" +) + +var ( + ErrModelNotFound = errors.New("model not found in registry") + ErrModelDisabled = errors.New("model is disabled") + ErrSwapInProgress = errors.New("model swap already in progress") + ErrAlreadyLoaded = errors.New("model already loaded") + ErrServerNotRunning = errors.New("llama-server not running") +) + +type Manager struct { + cfg *config.Config + db *gorm.DB + mu sync.RWMutex + state State + proc *exec.Cmd + done chan struct{} + url string +} + +func NewManager(cfg *config.Config, db *gorm.DB) *Manager { + return &Manager{ + cfg: cfg, + db: db, + done: make(chan struct{}), + url: cfg.LlamaServerURL(), + } +} + +func (m *Manager) GetDB() *gorm.DB { + return m.db +} + +func (m *Manager) Start() error { + m.mu.Lock() + defer m.mu.Unlock() + + slog.Info("llama manager starting", "url", m.url) + + // Load default model on startup + var model db.Model + if err := m.db.Where("is_default = ? AND is_enabled = ?", true, true).First(&model).Error; err == nil { + slog.Info("loading default model", "name", model.Name) + if err := m.loadModelInternal(&model); err != nil { + slog.Warn("failed to load default model", "error", err) + m.state.Status = StatusFailed + m.state.LastError = err.Error() + return nil + } + } + + // Start health check loop + go m.healthCheckLoop() + + return nil +} + +func (m *Manager) Stop() { + slog.Info("llama manager stopping") + close(m.done) + + m.mu.Lock() + defer m.mu.Unlock() + + if m.proc != nil && m.proc.Process != nil { + ctx, cancel := context.WithTimeout(context.Background(), m.cfg.LlamaServerStopTimeoutDuration()) + defer cancel() + + m.proc.SysProcAttr = &syscall.SysProcAttr{ + Setpgid: true, + } + + pgid, err := syscall.Getpgid(m.proc.Process.Pid) + if err == nil { + syscall.Kill(-pgid, syscall.SIGTERM) + } else { + m.proc.Process.Signal(syscall.SIGTERM) + } + + <-ctx.Done() + if m.proc.ProcessState == nil { + syscall.Kill(-pgid, syscall.SIGKILL) + } + } + + m.state = State{Status: StatusStopped} + slog.Info("llama manager stopped") +} + +func (m *Manager) IsReady() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.state.Status == StatusReady && m.proc != nil && m.proc.ProcessState != nil && !m.proc.ProcessState.Exited() +} + +func (m *Manager) Status() Status { + m.mu.RLock() + defer m.mu.RUnlock() + return m.state.Status +} + +func (m *Manager) CurrentModel() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.state.CurrentModel +} + +func (m *Manager) GetStatus() *State { + m.mu.RLock() + defer m.mu.RUnlock() + return &m.state +} + +func (m *Manager) LoadModel(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + + // Find model in DB + var model db.Model + if err := m.db.Where("name = ? AND is_enabled = ?", name, true).First(&model).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ErrModelNotFound + } + return err + } + + // Check current state + if m.state.Status == StatusLoading || m.state.Status == StatusSwapping { + if m.state.CurrentModel == name { + return nil // Already loading this model + } + return fmt.Errorf("swap in progress for %s, try again later", m.state.TargetModel) + } + + if m.state.CurrentModel == name && m.state.Status == StatusReady { + return nil // Already loaded + } + + return m.loadModelInternal(&model) +} + +func (m *Manager) loadModelInternal(model *db.Model) error { + isSwap := m.state.Status == StatusReady && m.state.CurrentModel != "" + m.state.Status = StatusSwapping + if !isSwap { + m.state.Status = StatusLoading + } + m.state.TargetModel = model.Name + m.state.LastError = "" + now := time.Now() + m.state.SwapStartedAt = &now + + slog.Info("loading model", "name", model.Name, "is_swap", isSwap) + + // Kill existing process + if m.proc != nil && m.proc.Process != nil { + m.terminateProcess() + } + + // Build command + cmd := m.buildCommand(model) + m.proc = cmd + + if err := cmd.Start(); err != nil { + m.state.Status = StatusFailed + m.state.LastError = err.Error() + return fmt.Errorf("failed to start llama-server: %w", err) + } + + slog.Info("llama-server started", "pid", cmd.Process.Pid) + m.state.PID = cmd.Process.Pid + + // Wait for server to be ready + if err := m.waitUntilReady(); err != nil { + m.state.Status = StatusFailed + m.state.LastError = err.Error() + return fmt.Errorf("model failed to start: %w", err) + } + + m.state.CurrentModel = model.Name + m.state.Status = StatusReady + m.state.TargetModel = "" + m.state.LoadedAt = &now + + // Update DB + m.db.Model(model).Updates(map[string]interface{}{ + "is_active": true, + "loaded_at": now, + }) + + slog.Info("model loaded successfully", "name", model.Name) + return nil +} + +func (m *Manager) buildCommand(model *db.Model) *exec.Cmd { + args := []string{ + "--model", model.ModelPath, + "--alias", model.Alias, + "--host", m.cfg.LlamaServerHost, + "--port", fmt.Sprintf("%d", m.cfg.LlamaServerPort), + "--ctx-size", fmt.Sprintf("%d", model.CtxSize), + "--n-gpu-layers", fmt.Sprintf("%d", model.NGPULayers), + } + + // Add extra args from JSON + extraArgs := ParseModelExtraArgs(model.ExtraArgs) + for k, v := range extraArgs { + if bv, ok := v.(bool); ok && bv { + args = append(args, "--"+k) + } else if v != nil { + args = append(args, "--"+k, fmt.Sprintf("%v", v)) + } + } + + cmd := exec.Command(m.cfg.LlamaServerBin, args...) + cmd.Stdout = io.Discard + cmd.Stderr = io.Discard + + // Set process group for clean kill + cmd.SysProcAttr = &syscall.SysProcAttr{ + Setpgid: true, + } + + return cmd +} + +func (m *Manager) terminateProcess() { + if m.proc == nil || m.proc.Process == nil { + return + } + + slog.Info("terminating llama-server", "pid", m.proc.Process.Pid) + + ctx, cancel := context.WithTimeout(context.Background(), m.cfg.LlamaServerStopTimeoutDuration()) + defer cancel() + + pgid, err := syscall.Getpgid(m.proc.Process.Pid) + if err == nil { + syscall.Kill(-pgid, syscall.SIGTERM) + } else { + m.proc.Process.Signal(syscall.SIGTERM) + } + + done := make(chan error, 1) + go func() { + done <- m.proc.Wait() + }() + + select { + case <-ctx.Done(): + if pgid, err := syscall.Getpgid(m.proc.Process.Pid); err == nil { + syscall.Kill(-pgid, syscall.SIGKILL) + } + case <-done: + } + + m.proc = nil +} + +func (m *Manager) waitUntilReady() error { + ctx, cancel := context.WithTimeout(context.Background(), m.cfg.LlamaServerStartupTimeoutDuration()) + defer cancel() + + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + if m.checkHealth() { + return nil + } + } + } +} + +func (m *Manager) checkHealth() bool { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, "GET", m.url+"/health", nil) + if err != nil { + return false + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + + return resp.StatusCode == http.StatusOK +} + +func (m *Manager) healthCheckLoop() { + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + + for { + select { + case <-m.done: + return + case <-ticker.C: + m.healthCheck() + } + } +} + +func (m *Manager) healthCheck() { + m.mu.RLock() + running := m.proc != nil && m.proc.Process != nil && m.proc.ProcessState != nil && !m.proc.ProcessState.Exited() + m.mu.RUnlock() + + if !running && m.state.Status == StatusReady { + m.mu.Lock() + m.state.Status = StatusFailed + m.state.LastError = "llama-server process died unexpectedly" + m.mu.Unlock() + slog.Error("llama-server process died", "current_model", m.state.CurrentModel) + } +} + +func (m *Manager) GetUsageStats() (totalRequests, totalTokens int64, avgLatencyMs float64) { + now := time.Now() + monthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + + var result struct { + TotalRequests int64 + TotalTokens int64 + AvgLatency float64 + } + + m.db.Model(&db.UsageLog{}). + Where("created_at >= ?", monthStart). + Select("COUNT(*) as total_requests, COALESCE(SUM(total_tokens), 0) as total_tokens, COALESCE(AVG(latency_ms), 0) as avg_latency"). + Scan(&result) + + return result.TotalRequests, result.TotalTokens, result.AvgLatency +} + +// ProxyRequest sends a request to the llama-server proxy +func (m *Manager) ProxyRequest(ctx context.Context, method, path string, body io.Reader, headers map[string]string) (*http.Response, error) { + if !m.IsReady() { + return nil, ErrServerNotRunning + } + + url := m.url + path + req, err := http.NewRequestWithContext(ctx, method, url, body) + if err != nil { + return nil, err + } + + req.Header.Set("Content-Type", "application/json") + for k, v := range headers { + req.Header.Set(k, v) + } + + return http.DefaultClient.Do(req) +} + +func ParseModelExtraArgs(extraArgs db.StringArray) map[string]interface{} { + if len(extraArgs) == 0 { + return nil + } + + // If it's a JSON string, parse it + if len(extraArgs) == 1 { + var result map[string]interface{} + if json.Unmarshal([]byte(extraArgs[0]), &result) == nil { + return result + } + } + + // Otherwise assume key=value pairs + result := make(map[string]interface{}) + for _, arg := range extraArgs { + parts := strings.SplitN(arg, "=", 2) + if len(parts) == 2 { + result[parts[0]] = parts[1] + } else { + result[arg] = true + } + } + return result +} diff --git a/internal/llama/proxy.go b/internal/llama/proxy.go new file mode 100644 index 0000000..2f32f3e --- /dev/null +++ b/internal/llama/proxy.go @@ -0,0 +1,247 @@ +package llama + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "strings" + "time" +) + +type Proxy struct { + manager *Manager + client *http.Client +} + +func NewProxy(manager *Manager) *Proxy { + return &Proxy{ + manager: manager, + client: &http.Client{ + Timeout: 60 * time.Second, + }, + } +} + +func (p *Proxy) Manager() *Manager { + return p.manager +} + +type ChatMessage struct { + Role string `json:"role"` + Content string `json:"content"` + Name string `json:"name,omitempty"` +} + +type ChatCompletionRequest struct { + Model string `json:"model"` + Messages []ChatMessage `json:"messages"` + Stream bool `json:"stream,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + TopP float64 `json:"top_p,omitempty"` +} + +type ChatCompletionResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []Choice `json:"choices"` + Usage Usage `json:"usage"` +} + +type Choice struct { + Index int `json:"index"` + Message ChatMessage `json:"message"` + FinishReason string `json:"finish_reason"` +} + +type Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +type StreamChoice struct { + Index int `json:"index"` + Delta ChatMessage `json:"delta"` + FinishReason string `json:"finish_reason,omitempty"` +} + +type StreamResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []StreamChoice `json:"choices"` +} + +// ChatCompletion calls llama-server and returns the response +func (p *Proxy) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) { + if !p.manager.IsReady() { + return nil, ErrServerNotRunning + } + + // Convert to llama-server format + llamaReq := map[string]interface{}{ + "model": req.Model, + "messages": req.Messages, + "stream": false, + } + + body, err := json.Marshal(llamaReq) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", p.manager.url+"/v1/chat/completions", bytes.NewReader(body)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := p.client.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("llama-server request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + bodyBytes, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("llama-server returned %d: %s", resp.StatusCode, string(bodyBytes)) + } + + var result ChatCompletionResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + + return &result, nil +} + +// ChatCompletionStream returns a channel of streaming responses +func (p *Proxy) ChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan *StreamResponse, <-chan error) { + stream := make(chan *StreamResponse, 100) + errCh := make(chan error, 1) + + if !p.manager.IsReady() { + errCh <- ErrServerNotRunning + close(stream) + return stream, errCh + } + + go func() { + defer close(stream) + defer close(errCh) + + llamaReq := map[string]interface{}{ + "model": req.Model, + "messages": req.Messages, + "stream": true, + } + + body, err := json.Marshal(llamaReq) + if err != nil { + errCh <- err + return + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", p.manager.url+"/v1/chat/completions", bytes.NewReader(body)) + if err != nil { + errCh <- err + return + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := p.client.Do(httpReq) + if err != nil { + errCh <- err + return + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + bodyBytes, _ := io.ReadAll(resp.Body) + errCh <- fmt.Errorf("llama-server returned %d: %s", resp.StatusCode, string(bodyBytes)) + return + } + + reader := bufio.NewReader(resp.Body) + for { + line, err := reader.ReadString('\n') + if err != nil { + if err != io.EOF { + errCh <- err + } + break + } + + line = strings.TrimSpace(line) + if line == "" || !strings.HasPrefix(line, "data: ") { + continue + } + + if line == "data: [DONE]" { + break + } + + data := strings.TrimPrefix(line, "data: ") + var streamResp StreamResponse + if err := json.Unmarshal([]byte(data), &streamResp); err != nil { + slog.Debug("failed to parse stream chunk", "error", err, "data", data) + continue + } + + select { + case stream <- &streamResp: + case <-ctx.Done(): + return + } + } + }() + + return stream, errCh +} + +// ModelsList returns available models from llama-server +func (p *Proxy) ModelsList(ctx context.Context) ([]string, error) { + if !p.manager.IsReady() { + return nil, ErrServerNotRunning + } + + req, err := http.NewRequestWithContext(ctx, "GET", p.manager.url+"/v1/models", nil) + if err != nil { + return nil, err + } + + resp, err := p.client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("llama-server returned %d", resp.StatusCode) + } + + var result struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, err + } + + models := make([]string, len(result.Data)) + for i, m := range result.Data { + models[i] = m.ID + } + return models, nil +} diff --git a/internal/llama/state.go b/internal/llama/state.go new file mode 100644 index 0000000..74dea3f --- /dev/null +++ b/internal/llama/state.go @@ -0,0 +1,34 @@ +package llama + +import "time" + +type Status string + +const ( + StatusStopped Status = "stopped" + StatusLoading Status = "loading" + StatusReady Status = "ready" + StatusSwapping Status = "swapping" + StatusFailed Status = "failed" +) + +type State struct { + CurrentModel string `json:"current_model"` + TargetModel string `json:"target_model,omitempty"` + Status Status `json:"status"` + PID int `json:"pid,omitempty"` + LoadedAt *time.Time `json:"loaded_at,omitempty"` + LastError string `json:"last_error,omitempty"` + SwapStartedAt *time.Time `json:"swap_started_at,omitempty"` + SwapInProgress bool `json:"swap_in_progress"` +} + +type ModelInfo struct { + Name string + ModelPath string + Alias string + CtxSize int + NGPULayers int + ExtraArgs map[string]interface{} + IsDefault bool +} diff --git a/internal/quota/service.go b/internal/quota/service.go new file mode 100644 index 0000000..2c5988b --- /dev/null +++ b/internal/quota/service.go @@ -0,0 +1,119 @@ +package quota + +import ( + "errors" + "log/slog" + "time" + + "github.com/google/uuid" + "gorm.io/gorm" + + "github.com/llamalink/llamalink/internal/db" +) + +var ( + ErrQuotaExceeded = errors.New("quota exceeded for this period") +) + +type Service struct { + db *gorm.DB +} + +func NewService(db *gorm.DB) *Service { + return &Service{db: db} +} + +func (s *Service) CheckQuota(apiKeyID uuid.UUID, modelScope string) (bool, string, error) { + now := time.Now().UTC() + + var quota db.Quota + err := s.db.Where( + "api_key_id = ? AND period_start <= ? AND period_end > ?", + apiKeyID, now, now, + ).First("a).Error + + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return true, "", nil // No quota configured + } + return false, "", err + } + + // Check model-specific scope if set + if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" { + if *quota.ModelScope != modelScope { + return true, "", nil // Quota doesn't apply to this model + } + } + + if quota.TokensUsed >= quota.TokensLimit { + return false, "quota exceeded", nil + } + + // Warning at 90% + if quota.TokensLimit > 0 { + usage := float64(quota.TokensUsed) / float64(quota.TokensLimit) + if usage >= 0.9 { + slog.Warn("quota usage warning", + "api_key_id", apiKeyID, + "usage_percent", usage*100, + "tokens_used", quota.TokensUsed, + "tokens_limit", quota.TokensLimit, + ) + } + } + + return true, "", nil +} + +func (s *Service) ConsumeQuota(apiKeyID uuid.UUID, modelScope string, tokens int) error { + now := time.Now().UTC() + periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + periodEnd := periodStart.AddDate(0, 1, 0) + + return s.db.Transaction(func(tx *gorm.DB) error { + var quota db.Quota + err := tx.Where( + "api_key_id = ? AND period_start = ? AND period_end = ?", + apiKeyID, periodStart, periodEnd, + ).First("a).Error + + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil // No quota configured + } + return err + } + + // Check model scope + if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" { + if *quota.ModelScope != modelScope { + return nil // Quota doesn't apply + } + } + + newUsed := quota.TokensUsed + tokens + if quota.TokensLimit > 0 && newUsed > quota.TokensLimit { + return ErrQuotaExceeded + } + + return tx.Model("a).Update("tokens_used", newUsed).Error + }) +} + +func (s *Service) GetUsage(apiKeyID uuid.UUID, periodStart, periodEnd time.Time) (used int, limit int, err error) { + var quota db.Quota + err = s.db.Where( + "api_key_id = ? AND period_start = ? AND period_end = ?", + apiKeyID, periodStart, periodEnd, + ).First("a).Error + + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return 0, 0, nil + } + return 0, 0, err + } + + return quota.TokensUsed, quota.TokensLimit, nil +} diff --git a/internal/quota/webhook.go b/internal/quota/webhook.go new file mode 100644 index 0000000..709a614 --- /dev/null +++ b/internal/quota/webhook.go @@ -0,0 +1,107 @@ +package quota + +import ( + "bytes" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "time" + + "github.com/google/uuid" + "gorm.io/gorm" + + "github.com/llamalink/llamalink/internal/db" +) + +type WebhookPayload struct { + Event string `json:"event"` + Timestamp time.Time `json:"timestamp"` + KeyID uuid.UUID `json:"key_id"` + KeyName string `json:"key_name"` + Message string `json:"message"` + Usage *struct { + Used int `json:"tokens_used"` + Limit int `json:"tokens_limit"` + } `json:"usage,omitempty"` +} + +type WebhookService struct { + db *gorm.DB + client *http.Client + webhookSecret string +} + +func NewWebhookService(db *gorm.DB) *WebhookService { + return &WebhookService{ + db: db, + client: &http.Client{Timeout: 10 * time.Second}, + } +} + +func (s *WebhookService) Dispatch(event string, key *db.ApiKey, message string, usage *struct{ Used, Limit int }) { + var webhooks []db.Webhook + s.db.Where("api_key_id = ? AND event = ? AND is_active = ?", key.ID, event, true).Find(&webhooks) + + for _, wh := range webhooks { + go s.send(wh, event, key, message, usage) + } +} + +func (s *WebhookService) send(wh db.Webhook, event string, key *db.ApiKey, message string, usage *struct{ Used, Limit int }) { + payload := WebhookPayload{ + Event: event, + Timestamp: time.Now().UTC(), + KeyID: key.ID, + KeyName: key.Name, + Message: message, + } + + if usage != nil { + payload.Usage = &struct { + Used int `json:"tokens_used"` + Limit int `json:"tokens_limit"` + }{Used: usage.Used, Limit: usage.Limit} + } + + body, err := json.Marshal(payload) + if err != nil { + slog.Error("failed to marshal webhook payload", "error", err) + return + } + + req, err := http.NewRequest("POST", wh.URL, bytes.NewReader(body)) + if err != nil { + slog.Error("failed to create webhook request", "error", err) + return + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-LlamaLink-Event", event) + req.Header.Set("X-LlamaLink-Timestamp", fmt.Sprintf("%d", time.Now().Unix())) + + if wh.Secret != nil && *wh.Secret != "" { + sig := s.sign(body, *wh.Secret) + req.Header.Set("X-LlamaLink-Signature", sig) + } + + resp, err := s.client.Do(req) + if err != nil { + slog.Error("webhook delivery failed", "url", wh.URL, "error", err) + return + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + slog.Warn("webhook returned error", "url", wh.URL, "status", resp.StatusCode) + } +} + +func (s *WebhookService) sign(body []byte, secret string) string { + h := hmac.New(sha256.New, []byte(secret)) + h.Write(body) + return "sha256=" + hex.EncodeToString(h.Sum(nil)) +} diff --git a/web/frontend/index.html b/web/frontend/index.html new file mode 100644 index 0000000..011b9e7 --- /dev/null +++ b/web/frontend/index.html @@ -0,0 +1,16 @@ + + + + + + + LlamaLink Admin + + + + + +
+ + + diff --git a/web/frontend/package.json b/web/frontend/package.json new file mode 100644 index 0000000..afd5261 --- /dev/null +++ b/web/frontend/package.json @@ -0,0 +1,38 @@ +{ + "name": "llamalink-admin", + "version": "0.1.0", + "private": true, + "type": "module", + "scripts": { + "dev": "vite", + "build": "vue-tsc -b && vite build", + "preview": "vite preview", + "lint": "eslint . --ext .vue,.js,.jsx,.cjs,.mjs,.ts,.tsx,.cts,.mts --fix" + }, + "dependencies": { + "vue": "^3.5.0", + "vue-router": "^4.5.0", + "pinia": "^2.3.0", + "axios": "^1.7.0", + "zod": "^3.23.0", + "@vueuse/core": "^12.0.0", + "chart.js": "^4.4.0", + "vue-chartjs": "^5.3.0", + "lucide-vue-next": "^0.460.0" + }, + "devDependencies": { + "@vitejs/plugin-vue": "^5.2.0", + "vite": "^6.0.0", + "vue-tsc": "^2.2.0", + "typescript": "~5.6.0", + "tailwindcss": "^3.4.0", + "autoprefixer": "^10.4.0", + "postcss": "^8.4.0", + "@nuxtjs/tailwindcss": "^8.0.0", + "@types/node": "^22.0.0", + "eslint": "^9.0.0", + "@typescript-eslint/eslint-plugin": "^8.0.0", + "@typescript-eslint/parser": "^8.0.0", + "eslint-plugin-vue": "^9.0.0" + } +} diff --git a/web/frontend/postcss.config.js b/web/frontend/postcss.config.js new file mode 100644 index 0000000..2e7af2b --- /dev/null +++ b/web/frontend/postcss.config.js @@ -0,0 +1,6 @@ +export default { + plugins: { + tailwindcss: {}, + autoprefixer: {}, + }, +} diff --git a/web/frontend/src/App.vue b/web/frontend/src/App.vue new file mode 100644 index 0000000..a78fb6f --- /dev/null +++ b/web/frontend/src/App.vue @@ -0,0 +1,7 @@ + + + diff --git a/web/frontend/src/env.d.ts b/web/frontend/src/env.d.ts new file mode 100644 index 0000000..323c78a --- /dev/null +++ b/web/frontend/src/env.d.ts @@ -0,0 +1,7 @@ +/// + +declare module '*.vue' { + import type { DefineComponent } from 'vue' + const component: DefineComponent<{}, {}, any> + export default component +} diff --git a/web/frontend/src/lib/api.ts b/web/frontend/src/lib/api.ts new file mode 100644 index 0000000..b7a759b --- /dev/null +++ b/web/frontend/src/lib/api.ts @@ -0,0 +1,20 @@ +import axios from 'axios' + +export const api = axios.create({ + baseURL: '/api', + timeout: 30000, + headers: { + 'Content-Type': 'application/json', + }, +}) + +api.interceptors.response.use( + (response) => response, + (error) => { + if (error.response?.status === 401) { + localStorage.removeItem('admin_token') + window.location.href = '/admin/login' + } + return Promise.reject(error) + } +) diff --git a/web/frontend/src/main.ts b/web/frontend/src/main.ts new file mode 100644 index 0000000..c4b3a5f --- /dev/null +++ b/web/frontend/src/main.ts @@ -0,0 +1,10 @@ +import { createApp } from 'vue' +import { createPinia } from 'pinia' +import router from './router' +import App from './App.vue' +import './style.css' + +const app = createApp(App) +app.use(createPinia()) +app.use(router) +app.mount('#app') diff --git a/web/frontend/src/router/index.ts b/web/frontend/src/router/index.ts new file mode 100644 index 0000000..83ce828 --- /dev/null +++ b/web/frontend/src/router/index.ts @@ -0,0 +1,61 @@ +import { createRouter, createWebHistory } from 'vue-router' +import { useAuthStore } from '@/stores/auth' + +const routes = [ + { + path: '/admin/login', + name: 'Login', + component: () => import('@/views/Login.vue'), + meta: { guest: true }, + }, + { + path: '/admin/', + component: () => import('@/views/Layout.vue'), + meta: { requiresAuth: true }, + children: [ + { + path: '', + name: 'Dashboard', + component: () => import('@/views/Dashboard.vue'), + }, + { + path: 'keys', + name: 'ApiKeys', + component: () => import('@/views/ApiKeys.vue'), + }, + { + path: 'models', + name: 'Models', + component: () => import('@/views/Models.vue'), + }, + { + path: 'usage', + name: 'Usage', + component: () => import('@/views/Usage.vue'), + }, + ], + }, + { + path: '/:pathMatch(.*)*', + redirect: '/admin/', + }, +] + +const router = createRouter({ + history: createWebHistory('/admin'), + routes, +}) + +router.beforeEach((to, from, next) => { + const authStore = useAuthStore() + + if (to.meta.requiresAuth && !authStore.isAuthenticated) { + next({ name: 'Login' }) + } else if (to.meta.guest && authStore.isAuthenticated) { + next({ name: 'Dashboard' }) + } else { + next() + } +}) + +export default router diff --git a/web/frontend/src/stores/auth.ts b/web/frontend/src/stores/auth.ts new file mode 100644 index 0000000..534d8fb --- /dev/null +++ b/web/frontend/src/stores/auth.ts @@ -0,0 +1,43 @@ +import { defineStore } from 'pinia' +import { ref, computed } from 'vue' +import { api } from '@/lib/api' + +export const useAuthStore = defineStore('auth', () => { + const token = ref(localStorage.getItem('admin_token')) + const loading = ref(false) + const error = ref(null) + + const isAuthenticated = computed(() => !!token.value) + + async function login(adminToken: string) { + loading.value = true + error.value = null + + try { + const response = await api.post('/api/v1/admin/login', { admin_token: adminToken }) + token.value = response.data.token + localStorage.setItem('admin_token', response.data.token) + api.defaults.headers.common['Authorization'] = `Bearer ${response.data.token}` + return true + } catch (err: any) { + error.value = err.response?.data?.error?.message || 'Login failed' + return false + } finally { + loading.value = false + } + } + + function logout() { + token.value = null + localStorage.removeItem('admin_token') + delete api.defaults.headers.common['Authorization'] + } + + function init() { + if (token.value) { + api.defaults.headers.common['Authorization'] = `Bearer ${token.value}` + } + } + + return { token, loading, error, isAuthenticated, login, logout, init } +}) diff --git a/web/frontend/src/style.css b/web/frontend/src/style.css new file mode 100644 index 0000000..d4c29d2 --- /dev/null +++ b/web/frontend/src/style.css @@ -0,0 +1,108 @@ +@tailwind base; +@tailwind components; +@tailwind utilities; + +* { + margin: 0; + padding: 0; + box-sizing: border-box; +} + +html { + font-family: 'Inter', system-ui, sans-serif; +} + +body { + background-color: #0d1117; + color: #e6edf3; + min-height: 100vh; +} + +code, pre { + font-family: 'JetBrains Mono', Consolas, monospace; +} + +@layer components { + .btn { + @apply px-4 py-2 rounded-lg font-medium transition-colors duration-200; + } + + .btn-primary { + @apply bg-primary hover:bg-primary-hover text-white; + } + + .btn-secondary { + @apply bg-surface border border-border hover:bg-border text-text; + } + + .btn-danger { + @apply bg-error hover:bg-red-600 text-white; + } + + .btn-sm { + @apply px-3 py-1.5 text-sm; + } + + .input { + @apply w-full px-3 py-2 bg-surface border border-border rounded-lg text-text placeholder-text-muted focus:outline-none focus:ring-2 focus:ring-primary focus:border-transparent; + } + + .card { + @apply bg-surface border border-border rounded-xl p-6; + } + + .badge { + @apply inline-flex items-center px-2.5 py-0.5 rounded-full text-xs font-medium; + } + + .badge-success { + @apply bg-success/15 text-success; + } + + .badge-warning { + @apply bg-warning/15 text-warning; + } + + .badge-error { + @apply bg-error/15 text-error; + } + + .badge-info { + @apply bg-primary/15 text-primary; + } + + .table { + @apply w-full text-left; + } + + .table th { + @apply px-4 py-3 text-xs font-medium text-text-muted uppercase tracking-wider border-b border-border; + } + + .table td { + @apply px-4 py-3 border-b border-border; + } + + .table tr:hover td { + @apply bg-surface; + } +} + +/* Scrollbar */ +::-webkit-scrollbar { + width: 8px; + height: 8px; +} + +::-webkit-scrollbar-track { + background: #0d1117; +} + +::-webkit-scrollbar-thumb { + background: #30363d; + border-radius: 4px; +} + +::-webkit-scrollbar-thumb:hover { + background: #484f58; +} diff --git a/web/frontend/src/views/ApiKeys.vue b/web/frontend/src/views/ApiKeys.vue new file mode 100644 index 0000000..a07d3fd --- /dev/null +++ b/web/frontend/src/views/ApiKeys.vue @@ -0,0 +1,172 @@ + + + diff --git a/web/frontend/src/views/Dashboard.vue b/web/frontend/src/views/Dashboard.vue new file mode 100644 index 0000000..17a0d91 --- /dev/null +++ b/web/frontend/src/views/Dashboard.vue @@ -0,0 +1,219 @@ + + + diff --git a/web/frontend/src/views/Layout.vue b/web/frontend/src/views/Layout.vue new file mode 100644 index 0000000..3c5dd3c --- /dev/null +++ b/web/frontend/src/views/Layout.vue @@ -0,0 +1,66 @@ + + + diff --git a/web/frontend/src/views/Login.vue b/web/frontend/src/views/Login.vue new file mode 100644 index 0000000..aa95834 --- /dev/null +++ b/web/frontend/src/views/Login.vue @@ -0,0 +1,70 @@ + + + diff --git a/web/frontend/src/views/Models.vue b/web/frontend/src/views/Models.vue new file mode 100644 index 0000000..5a8c9e6 --- /dev/null +++ b/web/frontend/src/views/Models.vue @@ -0,0 +1,168 @@ + + + diff --git a/web/frontend/src/views/Usage.vue b/web/frontend/src/views/Usage.vue new file mode 100644 index 0000000..513e84f --- /dev/null +++ b/web/frontend/src/views/Usage.vue @@ -0,0 +1,194 @@ + + + diff --git a/web/frontend/tailwind.config.js b/web/frontend/tailwind.config.js new file mode 100644 index 0000000..352702b --- /dev/null +++ b/web/frontend/tailwind.config.js @@ -0,0 +1,28 @@ +/** @type {import('tailwindcss').Config} */ +export default { + content: [ + "./index.html", + "./src/**/*.{vue,js,ts,jsx,tsx}", + ], + theme: { + extend: { + colors: { + background: '#0d1117', + surface: '#161b22', + border: '#30363d', + primary: '#6366f1', + 'primary-hover': '#818cf8', + success: '#22c55e', + warning: '#f59e0b', + error: '#ef4444', + text: '#e6edf3', + 'text-muted': '#8b949e', + }, + fontFamily: { + sans: ['Inter', 'system-ui', 'sans-serif'], + mono: ['JetBrains Mono', 'Consolas', 'monospace'], + }, + }, + }, + plugins: [], +} diff --git a/web/frontend/tsconfig.json b/web/frontend/tsconfig.json new file mode 100644 index 0000000..cee0024 --- /dev/null +++ b/web/frontend/tsconfig.json @@ -0,0 +1,25 @@ +{ + "compilerOptions": { + "target": "ES2022", + "useDefineForClassFields": true, + "module": "ESNext", + "lib": ["ES2022", "DOM", "DOM.Iterable"], + "skipLibCheck": true, + "moduleResolution": "bundler", + "allowImportingTsExtensions": true, + "resolveJsonModule": true, + "isolatedModules": true, + "noEmit": true, + "jsx": "preserve", + "strict": true, + "noUnusedLocals": true, + "noUnusedParameters": true, + "noFallthroughCasesInSwitch": true, + "paths": { + "@/*": ["./src/*"] + }, + "types": ["vite/client"] + }, + "include": ["src/**/*.ts", "src/**/*.tsx", "src/**/*.vue"], + "references": [{ "path": "./tsconfig.node.json" }] +} diff --git a/web/frontend/tsconfig.node.json b/web/frontend/tsconfig.node.json new file mode 100644 index 0000000..97ede7e --- /dev/null +++ b/web/frontend/tsconfig.node.json @@ -0,0 +1,11 @@ +{ + "compilerOptions": { + "composite": true, + "skipLibCheck": true, + "module": "ESNext", + "moduleResolution": "bundler", + "allowSyntheticDefaultImports": true, + "strict": true + }, + "include": ["vite.config.ts"] +} diff --git a/web/frontend/vite.config.ts b/web/frontend/vite.config.ts new file mode 100644 index 0000000..8bb31b7 --- /dev/null +++ b/web/frontend/vite.config.ts @@ -0,0 +1,29 @@ +import { defineConfig } from 'vite' +import vue from '@vitejs/plugin-vue' +import { resolve } from 'path' + +export default defineConfig({ + plugins: [vue()], + resolve: { + alias: { + '@': resolve(__dirname, 'src'), + }, + }, + server: { + port: 5173, + proxy: { + '/api': { + target: 'http://localhost:8000', + changeOrigin: true, + }, + '/admin': { + target: 'http://localhost:8000', + changeOrigin: true, + }, + }, + }, + build: { + outDir: '../dist', + emptyOutDir: true, + }, +})