bootstrap.go

  1// Package bootstrap builds the service clients shared by the kaiwari entry points:
  2// scenario load, self-hosted model services (unless disabled), one LLM client per
  3// prompt slot, and the prompt warmup. Every failure is returned as an error.
  4package bootstrap
  5
  6import (
  7	"context"
  8	"fmt"
  9	"net/http"
 10	"time"
 11
 12	"git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/config"
 13	"git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/llm"
 14	"git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/scenario"
 15	"git.theedgeofrage.com/TheEdgeOfRage/kaiwari/internal/services"
 16)
 17
 18const (
 19	warmupTimeout = 30 * time.Second
 20
 21	// llama-server slot assignments are fixed for this deployment; they are not
 22	// user-configurable. One slot per prompt family keeps each system prompt
 23	// cached in the server's unified KV pool; scratch questions and flashcards
 24	// share a slot since both are rare one-shot side tasks.
 25	gameSlot       = 0
 26	judgeSlot      = 1
 27	compactionSlot = 2
 28	sheetSlot      = 3
 29	scratchSlot    = 4
 30)
 31
 32// Services is the set of shared clients one kaiwari entry point builds at startup.
 33type Services struct {
 34	Cfg        *config.Config
 35	Brief      string
 36	HTTP       *http.Client
 37	Game       *llm.Client
 38	Judge      *llm.Client
 39	Compaction *llm.Client
 40	Sheet      *llm.Client
 41	Scratch    *llm.Client
 42
 43	manager *services.Manager
 44}
 45
 46// New loads the scenario, starts the self-hosted model services unless disabled,
 47// builds one LLM client per slot, and warms up the prompt slots. Any failure
 48// tears down whatever was started and is returned as an error. Cancelling ctx
 49// aborts startup and cleans up any spawned children.
 50func New(ctx context.Context, cfg *config.Config) (*Services, error) {
 51	sc, err := scenario.Load(cfg.Scenario)
 52	if err != nil {
 53		return nil, fmt.Errorf("load scenario: %w", err)
 54	}
 55
 56	c := *cfg
 57	var manager *services.Manager
 58	if !c.DisableModelLoading {
 59		manager = services.NewManager()
 60		if err := manager.Start(ctx); err != nil {
 61			return nil, err
 62		}
 63		// Managed children bind loopback; pin the endpoints regardless of URL flags.
 64		c.LLMConfig.BaseURL = "http://" + services.LLMListen
 65		c.AudioConfig.AudioBaseURL = "http://" + services.AudioListen
 66	}
 67
 68	hc := &http.Client{}
 69	gameClient := newLLMClient(&c, hc, gameSlot)
 70	judgeClient := newLLMClient(&c, hc, judgeSlot)
 71	compactionClient := newLLMClient(&c, hc, compactionSlot)
 72	sheetClient := newLLMClient(&c, hc, sheetSlot)
 73	scratchClient := newLLMClient(&c, hc, scratchSlot)
 74
 75	warmups := []struct {
 76		client *llm.Client
 77		prompt string
 78	}{
 79		{gameClient, llm.GameSystemPrompt(sc.Brief)},
 80		{judgeClient, llm.JudgeSystemPrompt()},
 81		{compactionClient, llm.CompactionPrompt()},
 82		{sheetClient, llm.SheetSystemPrompt()},
 83	}
 84	warmCtx, cancel := context.WithTimeout(ctx, warmupTimeout)
 85	defer cancel()
 86	var warmErr error
 87	for _, w := range warmups {
 88		if warmErr = w.client.Warmup(warmCtx, []string{w.prompt}); warmErr != nil {
 89			break
 90		}
 91	}
 92	if warmErr != nil {
 93		if manager != nil {
 94			manager.Stop()
 95		}
 96		return nil, fmt.Errorf("llm warmup: %w", warmErr)
 97	}
 98
 99	return &Services{
100		Cfg:        &c,
101		Brief:      sc.Brief,
102		HTTP:       hc,
103		Game:       gameClient,
104		Judge:      judgeClient,
105		Compaction: compactionClient,
106		Sheet:      sheetClient,
107		Scratch:    scratchClient,
108		manager:    manager,
109	}, nil
110}
111
112// Close stops the spawned model services if any were started. Safe to call when
113// none were and safe to call multiple times.
114func (s *Services) Close() {
115	if s.manager != nil {
116		s.manager.Stop()
117	}
118}
119
120func newLLMClient(cfg *config.Config, hc *http.Client, slot int) *llm.Client {
121	c := llm.NewClient(cfg, hc)
122	c.Slot = slot
123	return c
124}