llmgw/internal/router/chat.go

package router

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"slices"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
	"gitea.mynodes.xyz/ilker/llmgw/internal/web"
)

// registerChat mounts the in-browser chat for signed-in people.
//
// It exists so someone can try a model without setting up a client, and — more
// often — so they can tell whether a problem is the gateway or their code.
//
// The browser never sees an API key. The page carries a session, the server
// turns that into the same identity an API key would produce, and the request
// goes down the ordinary path: admission, model allowlist, rate limits, budget,
// accounting. Anything else would make this page a hole in every limit the rest
// of the gateway enforces.
func (r *Router) registerChat(mux *http.ServeMux, pages *web.AccountPages) {
	mux.HandleFunc("GET /chat", pages.Guard(func(w http.ResponseWriter, req *http.Request, user store.User) {
		contexts, _ := r.chatContexts(req.Context(), user)
		current, err := r.pickContext(req.Context(), user, req.URL.Query().Get("context"))
		if err != nil {
			current = ""
		}
		if current == "" && !r.personalOff {
			current = personalContext
		}
		web.Chat(user, r.modelsFor(req, user), contexts, current, version.Version, r.chatKnowledge(req, user, current))(w, req)
	}))
	mux.HandleFunc("POST /chat/send", pages.Guard(r.chatSend))
	// Where a model stands while it loads, and a way to ask for it without
	// holding a request open until it is up.
	mux.HandleFunc("GET /chat/model-status", pages.Guard(r.handleChatStatus))
	mux.HandleFunc("POST /chat/wake", pages.Guard(r.handleChatWake))
	// Behind the same guard as the rest of the page: it takes a POST body and
	// returns a file, so it wants a session and the CSRF check for the same
	// reasons /chat/send does.
	mux.HandleFunc("POST /chat/download", pages.Guard(
		func(w http.ResponseWriter, req *http.Request, _ store.User) {
			web.CodeArchive(r.cfg.MaxBodyBytes)(w, req)
		}))
	// A finished reply, rendered: markdown the model wrote, made safe.
	mux.HandleFunc("POST /chat/render", pages.Guard(
		func(w http.ResponseWriter, req *http.Request, _ store.User) {
			web.ReplyRenderer(r.cfg.MaxBodyBytes)(w, req)
		}))
}

// chatSend dispatches one turn on behalf of a signed-in person.
func (r *Router) chatSend(w http.ResponseWriter, req *http.Request, user store.User) {
	body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
	if err != nil {
		http.Error(w, "could not read the request", http.StatusBadRequest)
		return
	}

	var in struct {
		Model    string            `json:"model"`
		Messages []json.RawMessage `json:"messages"`
		// Context is what this turn spends against: a project id, or
		// "personal". Checked here, not trusted from the page.
		Context string `json:"context"`
		// Temperature is the page's one sampling setting, held to the range
		// every engine takes.
		Temperature *float64 `json:"temperature"`
		// Knowledge is a collection to answer from, "none", or "" for the
		// default; whether it may be read is decided as for a key.
		Knowledge string `json:"knowledge"`
	}
	if err := json.Unmarshal(body, &in); err != nil || in.Model == "" || len(in.Messages) == 0 {
		http.Error(w, "a model and at least one message are required", http.StatusBadRequest)
		return
	}
	principal, err := r.principalFor(req, user, in.Context)
	if err != nil {
		http.Error(w, err.Error(), http.StatusForbidden)
		return
	}

	// Rebuilt rather than forwarded, so the page cannot set fields it has no
	// business setting — a stream:false that would hang the reader, a
	// max_tokens large enough to occupy a replica for an hour.
	_, model := r.resolve(in.Model)
	req2 := map[string]any{
		"model":      in.Model,
		"messages":   in.Messages,
		"stream":     true,
		"max_tokens": chatTokenCap(model),
		// So the page can say how many tokens the answer took, and how fast
		// they came, from the engine's own count.
		"stream_options": map[string]any{"include_usage": true},
	}
	if t := in.Temperature; t != nil {
		req2["temperature"] = max(0, min(*t, 2))
	}
	upstream, err := json.Marshal(req2)
	if err != nil {
		http.Error(w, "could not build the request", http.StatusInternalServerError)
		return
	}

	next := req.Clone(auth.WithPrincipal(req.Context(), principal))
	next.Body = io.NopCloser(bytes.NewReader(upstream))
	next.ContentLength = int64(len(upstream))
	next.Header.Set("Content-Type", "application/json")
	next.Header.Del(knowledgeHeader)
	if k := strings.TrimSpace(in.Knowledge); k != "" {
		next.Header.Set(knowledgeHeader, k)
	}

	// The same handler the API serves, so there is exactly one request path and
	// no second place for its rules to drift.
	r.handleInference("/v1/chat/completions")(w, next)
}

// chatMaxTokens bounds one turn from the page.
//
// A playground should not be the thing that ties up a replica for ten minutes,
// and someone who wants a longer answer has an API key for that.
//
// 512 was chosen when every model here was a 0.5B answering questions in a
// paragraph. It is about a screenful of code, so a coding model cut off
// mid-function on the first real request — and the page said nothing, because
// it never read finish_reason. The cap being too low was a nuisance; the reply
// ending silently was the actual defect, and both are fixed.
const chatMaxTokens = 2048

// chatTokenCap is that bound, lowered to whatever the model allows.
//
// Sending the constant regardless made the gateway refuse its own page: a model
// declaring max_output 1024 answered every turn from /chat with "max_tokens is
// 2048, but this model allows at most 1024", which is a true sentence about a
// number the person never chose and cannot change. The ceiling is the page's to
// impose on itself, and the model's limit is not negotiable, so the smaller of
// the two is the only answer that is right in both directions.
func chatTokenCap(m catalog.Model) int {
	if m.MaxOutput > 0 && m.MaxOutput < chatMaxTokens {
		return m.MaxOutput
	}
	return chatMaxTokens
}

// personalContext is how the page names spending as oneself, outside any
// project.
const personalContext = "personal"

// chatContexts is what a person may spend against on this page: the projects
// they can use, and themselves when personal use is allowed.
func (r *Router) chatContexts(ctx context.Context, user store.User) ([]web.ChatContext, error) {
	var out []web.ChatContext
	projects, err := r.store.ProjectsFor(ctx, user.ID)
	if err != nil {
		return nil, err
	}
	for _, p := range projects {
		out = append(out, web.ChatContext{Value: p.ID, Label: p.Path})
	}
	if !r.personalOff {
		out = append(out, web.ChatContext{Value: personalContext, Label: "Personal"})
	}
	return out, nil
}

// pickContext resolves what a request asked to spend against, checking the
// person may; with nothing asked, their first project, else themselves.
// Returns the project id, empty for personal.
func (r *Router) pickContext(ctx context.Context, user store.User, asked string) (string, error) {
	switch {
	case asked == personalContext:
		if r.personalOff {
			return "", errors.New("personal use is off on this gateway")
		}
		return "", nil
	case asked != "":
		ok, err := r.store.CanUseProject(ctx, user.ID, asked)
		if err != nil {
			return "", err
		}
		if !ok {
			return "", errors.New("you are not a member of that project")
		}
		return asked, nil
	}
	projects, err := r.store.ProjectsFor(ctx, user.ID)
	if err != nil {
		return "", err
	}
	if len(projects) > 0 {
		return projects[0].ID, nil
	}
	if r.personalOff {
		return "", errors.New("you are in no project, and personal use is off on this gateway")
	}
	return "", nil
}

// principalFor builds the identity a signed-in person acts under, spending
// against a project they may use or, with none, themselves.
//
// The key id is synthetic and marked as a session, because there is no key: it
// keeps usage rows attributable and rate limits per person, while making it
// obvious in the usage table that these requests came from the page rather than
// from a credential somebody holds.
func (r *Router) principalFor(req *http.Request, user store.User, asked string) (auth.Principal, error) {
	if r.store == nil {
		return auth.Principal{}, errors.New("no store")
	}
	project, err := r.pickContext(req.Context(), user, asked)
	if err != nil {
		return auth.Principal{}, err
	}
	eff, err := r.store.ResolveFor(req.Context(), "session:"+user.ID, user.ID, project)
	if err != nil {
		return auth.Principal{}, err
	}
	return auth.Principal{
		KeyID:          "session:" + user.ID,
		KeyPrefix:      "session",
		UserID:         user.ID,
		ProjectID:      eff.ProjectID,
		TeamID:         eff.TeamID,
		OrgID:          eff.OrgID,
		Band:           eff.Band,
		RoutingProfile: eff.RoutingProfile,
		Limits:         eff,
	}, nil
}

// modelsFor lists what this person may actually ask for.
//
// Served rather than catalogued: offering a model with no replica produces a
// 503 the reader cannot interpret, and the allowlist is applied so the picker
// never shows a model the request would refuse.
//
// One entry per model, under its catalog name. An endpoint advertises every
// name a client may address it by — the model's own and all of its aliases,
// because that is what the router matches against — so taking those strings as
// they come listed the same model once for each name it has. A picker offering
// "coder", "coder-7b" and "qwen2.5-coder-7b-instruct" as three choices invites
// somebody to compare them.
//
// Real models are ordered ahead of simulated ones so the browser's default
// selection — the first option — is always a model that answers. Sorted purely
// by name, "big-model" came first and every visitor who pressed Send without
// touching the picker got placeholder text.
func (r *Router) modelsFor(req *http.Request, user store.User) []web.ChatModel {
	principal, err := r.principalFor(req, user, req.URL.Query().Get("context"))
	if err != nil {
		return nil
	}
	seen := map[string]struct{}{}
	var out []web.ChatModel
	// What each model's replicas compute on, across all of them: one model
	// can run on a card and on CPUs at once.
	hardware := map[string][]string{}
	for _, s := range r.reg.Snapshots("") {
		hw := hardwareOf(s.Pool)
		for _, m := range s.Models {
			name, _ := r.resolve(m)
			if hw != "" && !slices.Contains(hardware[name], hw) {
				hardware[name] = append(hardware[name], hw)
			}
		}
	}
	for _, s := range r.reg.Snapshots("") {
		for _, m := range s.Models {
			// Resolved rather than taken as sent: an alias and the name it
			// stands for are one model, and the catalog is what says so.
			name, entry := r.resolve(m)
			if _, dup := seen[name]; dup {
				continue
			}
			seen[name] = struct{}{}
			if r.permits(principal, name) {
				out = append(out, web.ChatModel{Name: name, Simulated: entry.Simulated(), Hardware: hardware[name]})
			}
		}
	}
	// A model that loads on demand has no endpoint until somebody asks for it,
	// and sending it a message is how somebody asks. Listing only what already
	// runs left such a model unreachable from this page for good.
	for name, entry := range r.modelMap() {
		if name != entry.Name || !entry.Scaling.OnDemand() {
			continue // an alias, or a model that is placed without being asked
		}
		if _, up := seen[name]; up {
			continue
		}
		seen[name] = struct{}{}
		if r.permits(principal, name) {
			out = append(out, web.ChatModel{
				Name: name, Simulated: entry.Simulated(), OnDemand: true,
			})
		}
	}
	slices.SortFunc(out, func(a, b web.ChatModel) int {
		if a.Simulated != b.Simulated {
			if b.Simulated {
				return -1
			}
			return 1
		}
		return strings.Compare(a.Name, b.Name)
	})
	return out
}

// chatKnowledge is the collections the page offers in a context, and the
// default that answers when none is chosen.
func (r *Router) chatKnowledge(req *http.Request, user store.User, context string) web.ChatKnowledge {
	var out web.ChatKnowledge
	if r.store == nil || r.store.KnowledgeAvailable() != nil {
		return out
	}
	p, err := r.principalFor(req, user, context)
	if err != nil {
		return out
	}
	e := store.Effective{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID, OrgID: p.OrgID}
	cs, err := r.store.CollectionsFor(req.Context(), e, false)
	if err != nil {
		return out
	}
	names := map[string]string{}
	for _, c := range cs {
		out.Collections = append(out.Collections, web.ChatCollection{ID: c.ID, Label: c.Name})
		names[c.ID] = c.Name
	}
	if def, err := r.store.KBDefaultFor(req.Context(), e); err == nil {
		var shown []string
		for _, id := range def {
			if n, ok := names[id]; ok {
				shown = append(shown, n)
			}
		}
		out.Default = strings.Join(shown, ", ")
	}
	return out
}