ai-service/cmd/worker/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"ai-service/internal/config"
	"ai-service/internal/llm"
	"ai-service/internal/migrate"
	"ai-service/internal/store"
	"ai-service/internal/transcription"
	"ai-service/internal/worker"
)

func main() {
	cfg := config.Load()
	slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil)))

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	db, err := store.Open(ctx, cfg.DatabaseURL)
	if err != nil {
		slog.Error("db_open_failed", "error", err)
		os.Exit(1)
	}
	defer db.Close()

	if cfg.MigrateOnStart {
		if err := migrate.Up(ctx, db); err != nil {
			slog.Error("migrate_failed", "error", err)
			os.Exit(1)
		}
	}

	if cfg.LLMBaseURL == "" {
		slog.Error("llm_not_configured")
		os.Exit(1)
	}

	llmClient := llm.New(cfg.LLMBaseURL, cfg.LLMAPIKey, cfg.LLMModel, cfg.LLMTimeout)
	transcriber := transcription.NewWithOptions(transcription.Options{
		VoxtralBaseURL: cfg.VoxtralBaseURL,
		VoxtralAPIKey:  cfg.VoxtralAPIKey,
		VoxtralModel:   cfg.VoxtralModel,
		VoxtralTimeout: cfg.VoxtralTimeout,
		AudioLLMPrompt: cfg.AudioLLMPrompt,
	})
	w := worker.New(db, llmClient, transcriber, cfg.WorkerID, cfg.LLMModel, cfg.WorkerTaskTypes, cfg.WorkerModelProfiles, cfg.WorkerPollInterval, cfg.WorkerLeaseTimeout, cfg.WorkerClaimLimit)
	healthSrv := startHealthServer(ctx, db, cfg)

	slog.Info("ai_worker_started",
		"worker_id", cfg.WorkerID,
		"model", cfg.LLMModel,
		"transcription_enabled", transcriber != nil,
		"transcription_provider", "voxtral-small",
		"transcription_model", cfg.VoxtralModel,
		"task_types", cfg.WorkerTaskTypes,
		"model_profiles", cfg.WorkerModelProfiles,
		"poll_interval", cfg.WorkerPollInterval.String(),
		"lease_timeout", cfg.WorkerLeaseTimeout.String(),
		"claim_limit", cfg.WorkerClaimLimit,
	)
	w.Run(ctx)

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = healthSrv.Shutdown(shutdownCtx)
}

type workerHealth struct {
	store *store.Store
	cfg   config.Config
}

func startHealthServer(ctx context.Context, db *store.Store, cfg config.Config) *http.Server {
	srv := &http.Server{
		Addr:              fmt.Sprintf("%s:%d", cfg.WorkerHTTPHost, cfg.WorkerHTTPPort),
		Handler:           workerHealth{store: db, cfg: cfg},
		ReadHeaderTimeout: 5 * time.Second,
	}
	go func() {
		slog.Info("ai_worker_health_started", "addr", srv.Addr)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			slog.Error("ai_worker_health_failed", "error", err)
		}
	}()
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = srv.Shutdown(shutdownCtx)
	}()
	return srv
}

func (h workerHealth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	path := strings.TrimSuffix(r.URL.Path, "/")
	if path == "" {
		path = "/"
	}
	switch {
	case r.Method == http.MethodGet && path == "/healthz":
		writeWorkerJSON(w, http.StatusOK, map[string]any{
			"status":    "ok",
			"worker_id": h.cfg.WorkerID,
		})
	case r.Method == http.MethodGet && path == "/readyz":
		ctx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
		defer cancel()
		if err := h.store.Ping(ctx); err != nil {
			writeWorkerJSON(w, http.StatusServiceUnavailable, map[string]any{
				"status": "not_ready",
				"error":  err.Error(),
			})
			return
		}
		writeWorkerJSON(w, http.StatusOK, map[string]any{
			"status":    "ready",
			"worker_id": h.cfg.WorkerID,
		})
	case r.Method == http.MethodGet && path == "/worker/status":
		writeWorkerJSON(w, http.StatusOK, map[string]any{
			"status":                 "running",
			"worker_id":              h.cfg.WorkerID,
			"task_types":             h.cfg.WorkerTaskTypes,
			"model_profiles":         h.cfg.WorkerModelProfiles,
			"transcription_provider": "voxtral-small",
			"transcription_model":    h.cfg.VoxtralModel,
			"claim_limit":            h.cfg.WorkerClaimLimit,
			"poll_interval":          h.cfg.WorkerPollInterval.String(),
			"lease_timeout":          h.cfg.WorkerLeaseTimeout.String(),
		})
	default:
		writeWorkerJSON(w, http.StatusNotFound, map[string]any{"error": "not found"})
	}
}

func writeWorkerJSON(w http.ResponseWriter, code int, body any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	if err := json.NewEncoder(w).Encode(body); err != nil {
		slog.Warn("worker_health_write_failed", "error", err)
	}
}