odidere/internal/service/inference.go

package service

import (
	"database/sql"
	"encoding/json/v2"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"code.chimeric.al/chimerical/odidere/internal/conversation"
	"code.chimeric.al/chimerical/odidere/internal/llm"
)

func (svc *Service) inference(w http.ResponseWriter, r *http.Request) {
	var (
		ctx = r.Context()
		log = ctx.Value(logKey).(*slog.Logger)
		id  = r.PathValue("id")
		req struct {
			Message              *llm.Message  `json:"message"`
			Provider             string        `json:"provider"`
			Model                string        `json:"model"`
			MaxIterations        int           `json:"max_iterations"`
			ThinkingBudgetTokens *int          `json:"thinking_budget_tokens"`
			Timeout              time.Duration `json:"timeout"`
			Stream               *bool         `json:"stream"`
		}
	)

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(
			w,
			"streaming not supported",
			http.StatusInternalServerError,
		)
		return
	}

	r.Body = http.MaxBytesReader(w, r.Body, 32<<20)
	if err := json.UnmarshalRead(r.Body, &req); err != nil {
		log.ErrorContext(
			ctx, "failed to decode request",
			slog.Any("error", err),
		)
		http.Error(w, "invalid request", http.StatusBadRequest)
		return
	}

	provider := req.Provider
	if provider == "" {
		provider = svc.cfg.DefaultModel.Provider
	}
	client := svc.llms[provider]
	if client == nil {
		log.ErrorContext(
			ctx, "unknown provider",
			slog.String("provider", provider),
		)
		http.Error(
			w,
			fmt.Sprintf("unknown provider: %q", provider),
			http.StatusBadRequest,
		)
		return
	}

	model := req.Model
	if model == "" {
		model = svc.cfg.DefaultModel.Model
	}
	if allowed, ok := svc.allowedModels[provider]; ok {
		base, _, _ := strings.Cut(model, ":")
		if _, ok := allowed[base]; !ok {
			log.ErrorContext(
				ctx, "model not in allowed list",
				slog.String("provider", provider),
				slog.String("model", model),
			)
			http.Error(
				w,
				fmt.Sprintf(
					"model %q not allowed for provider %q",
					model, provider,
				),
				http.StatusBadRequest,
			)
			return
		}
	}

	if _, err := svc.dbr.SelectConversationIdUserId(
		ctx, id, "default",
	); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			http.NotFound(w, r)
			return
		}
		log.ErrorContext(
			ctx, "failed to load conversation",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}
	c, err := svc.getConversation(ctx, id)
	if errors.Is(err, sql.ErrNoRows) {
		http.NotFound(w, r)
		return
	}
	if err != nil {
		log.ErrorContext(
			ctx, "failed to get conversation handle",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}

	stream := true
	if req.Stream != nil {
		stream = *req.Stream
	}
	inf, err := c.Reserve(svc.ctx, client, conversation.InferenceRequest{
		Message:              req.Message,
		Model:                model,
		MaxIterations:        req.MaxIterations,
		ThinkingBudgetTokens: req.ThinkingBudgetTokens,
		Timeout:              req.Timeout,
		Stream:               stream,
	})
	if errors.Is(err, conversation.ErrBusy) {
		http.Error(
			w,
			http.StatusText(http.StatusConflict),
			http.StatusConflict,
		)
		return
	}
	if err != nil {
		log.ErrorContext(
			ctx, "failed to reserve inference",
			slog.Any("error", err),
		)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Register the background run before returning so Service.Run waits
	// for it after HTTP shutdown completes.
	svc.wg.Go(func() {
		if err := inf.Run(); err != nil {
			log.ErrorContext(
				ctx, "background inference failed",
				slog.Any("error", err),
			)
		}
	})

	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Content-Type", "text/event-stream")
	for e := range inf.Events(ctx) {
		data, err := json.Marshal(e)
		if err != nil {
			log.ErrorContext(
				ctx, "failed to marshal SSE event",
				slog.Any("error", err),
			)
			return
		}
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", e.Type(), data)
		flusher.Flush()
	}
}

func (svc *Service) stop(w http.ResponseWriter, r *http.Request) {
	var (
		ctx = r.Context()
		log = ctx.Value(logKey).(*slog.Logger)
		id  = r.PathValue("id")
	)

	if _, err := svc.dbr.SelectConversationIdUserId(
		ctx, id, "default",
	); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			http.NotFound(w, r)
			return
		}
		log.ErrorContext(
			ctx, "failed to load conversation",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}
	c, err := svc.getConversation(ctx, id)
	if errors.Is(err, sql.ErrNoRows) {
		http.NotFound(w, r)
		return
	}
	if err != nil {
		log.ErrorContext(
			ctx, "failed to get conversation handle",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}

	c.Stop()
	w.WriteHeader(http.StatusNoContent)
}

func (svc *Service) subscribe(w http.ResponseWriter, r *http.Request) {
	var (
		ctx = r.Context()
		log = ctx.Value(logKey).(*slog.Logger)
		id  = r.PathValue("id")
		req struct {
			Manifest []struct {
				Id        string    `json:"id"`
				UpdatedAt time.Time `json:"updated_at"`
			} `json:"manifest"`
		}
	)

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(
			w,
			"streaming not supported",
			http.StatusInternalServerError,
		)
		return
	}

	r.Body = http.MaxBytesReader(w, r.Body, 32<<20)
	if err := json.UnmarshalRead(r.Body, &req); err != nil {
		log.ErrorContext(
			ctx, "failed to decode request",
			slog.Any("error", err),
		)
		http.Error(w, "invalid request", http.StatusBadRequest)
		return
	}

	if _, err := svc.dbr.SelectConversationIdUserId(
		ctx, id, "default",
	); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			http.NotFound(w, r)
			return
		}
		log.ErrorContext(
			ctx, "failed to load conversation",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}
	c, err := svc.getConversation(ctx, id)
	if errors.Is(err, sql.ErrNoRows) {
		http.NotFound(w, r)
		return
	}
	if err != nil {
		log.ErrorContext(
			ctx, "failed to get conversation handle",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}

	manifest := make([]conversation.MessageManifest, len(req.Manifest))
	for i, m := range req.Manifest {
		manifest[i] = conversation.MessageManifest{
			Id:        m.Id,
			UpdatedAt: m.UpdatedAt,
		}
	}
	events, err := c.Subscribe(ctx, manifest)
	if err != nil {
		log.ErrorContext(
			ctx, "failed to subscribe to conversation",
			slog.Any("error", err),
		)
		http.Error(
			w,
			http.StatusText(http.StatusInternalServerError),
			http.StatusInternalServerError,
		)
		return
	}

	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Content-Type", "text/event-stream")
	for e := range events {
		data, err := json.Marshal(e)
		if err != nil {
			log.ErrorContext(
				ctx, "failed to marshal SSE event",
				slog.Any("error", err),
			)
			return
		}
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", e.Type(), data)
		flusher.Flush()
	}
}