llmgw/internal/cli/ctl_request.go

package cli

import (
	"encoding/json"
	"fmt"
	"net/url"
	"slices"
	"strings"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func newRequestCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var (
		body   bool
		replay string
	)
	c := &cobra.Command{
		Use:   "request ID",
		Short: "What one request went through: replica, queue, first token, tokens, guardrails",
		Long: "Every response carries its id in X-Request-Id, and a client may send its\n" +
			"own. Given one, this says which replica answered, how long admission held\n" +
			"it, how long the first token took and what it cost — the answer to \"why\n" +
			"was this slow\". The access log has the same fields under \"request\".\n\n" +
			"For a project that keeps its requests (`llmgw keys requests`), --body\n" +
			"shows what the model was sent and what came back, and --replay MODEL\n" +
			"sends the same request to another model, past the limits and the\n" +
			"guardrails, and shows its answer beside the first.",
		Example: "  llmgw ctl request req_3f2a…\n  llmgw ctl request req_3f2a… --body\n  llmgw ctl request req_3f2a… --replay qwen2.5-14b-instruct",
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if replay != "" {
				return replayRequest(cmd, client, args[0], replay, *asJSON)
			}
			if body {
				return showRequestBody(cmd, client, args[0], *asJSON)
			}
			var v adminapi.RequestView
			if err := client().get(cmd.Context(), "/api/v1/requests/"+url.PathEscape(args[0]), &v); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
			}
			model := v.Model
			if v.ServedModel != "" && v.ServedModel != v.Model {
				model += " (served by " + v.ServedModel + ")"
			}
			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "request   %s at %s\n", v.RequestID, v.At.UTC().Format("2006-01-02 15:04:05Z"))
			fmt.Fprintf(out, "model     %s\n", model)
			fmt.Fprintf(out, "replica   %s\n", or(v.Replica, "- (none answered)"))
			fmt.Fprintf(out, "status    %d\n", v.Status)
			fmt.Fprintf(out, "time      %d ms total: %d queued, %d to the first token\n", v.LatencyMs, v.QueueMs, v.TTFTMs)
			fmt.Fprintf(out, "tokens    %d in, %d out\n", v.PromptTokens, v.CompletionTokens)
			if v.Images > 0 {
				fmt.Fprintf(out, "images    %d\n", v.Images)
			}
			if v.Meshes > 0 {
				fmt.Fprintf(out, "meshes    %d\n", v.Meshes)
			}
			if v.CacheHit {
				fmt.Fprintf(out, "cache     answered from the project's answer cache: nothing generated, nothing charged\n")
			}
			fmt.Fprintf(out, "who       key %s, project %s%s\n", v.KeyID, v.ProjectID, prefixed(", user ", v.UserID))
			if v.EndUser != "" {
				fmt.Fprintf(out, "end user  %s (the client's own name for them)\n", v.EndUser)
			}
			if len(v.Tags) > 0 {
				keys := make([]string, 0, len(v.Tags))
				for k := range v.Tags {
					keys = append(keys, k)
				}
				slices.Sort(keys)
				var kv []string
				for _, k := range keys {
					kv = append(kv, k+"="+v.Tags[k])
				}
				fmt.Fprintf(out, "tags      %s\n", strings.Join(kv, ", "))
			}
			if v.GuardAction != "" {
				what := or(v.GuardCategory, "-")
				if v.GuardRule != "" {
					what += ": " + v.GuardRule
				}
				fmt.Fprintf(out, "guard     %s by policy %s (%s), %d ms\n", v.GuardAction, v.GuardPolicy, what, v.GuardMs)
			} else if v.GuardMs > 0 {
				fmt.Fprintf(out, "guard     allowed, %d ms\n", v.GuardMs)
			}
			return nil
		},
	}
	c.Flags().BoolVar(&body, "body", false, "show the request and its answer, for a project that keeps them")
	c.Flags().StringVar(&replay, "replay", "", "send the kept request to this model and show both answers")
	return c
}

func showRequestBody(cmd *cobra.Command, client func() *ctlClient, id string, asJSON bool) error {
	var b store.RequestBody
	if err := client().get(cmd.Context(), "/api/v1/requests/"+url.PathEscape(id)+"/body", &b); err != nil {
		return err
	}
	if asJSON {
		return json.NewEncoder(cmd.OutOrStdout()).Encode(b)
	}
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "request   %s at %s, %s to %s (served by %s), status %d\n", b.RequestID,
		b.At.UTC().Format("2006-01-02 15:04:05Z"), b.Path, b.Model, b.ServedModel, b.Status)
	fmt.Fprintf(out, "kept      until %s\n\n", b.ExpiresAt.UTC().Format("2006-01-02 15:04Z"))
	fmt.Fprintln(out, "--- sent")
	fmt.Fprintln(out, indentJSON(b.Request))
	fmt.Fprintln(out, "--- answered")
	fmt.Fprintln(out, or(b.Response, "(no text: an error, or not a completion)"))
	if b.Truncated {
		fmt.Fprintln(out, "(longer than what is kept; cut here)")
	}
	return nil
}

func replayRequest(cmd *cobra.Command, client func() *ctlClient, id, model string, asJSON bool) error {
	var v adminapi.ReplayView
	body, _ := json.Marshal(map[string]string{"model": model})
	if err := client().post(cmd.Context(), "/api/v1/requests/"+url.PathEscape(id)+"/replay", body, &v); err != nil {
		return err
	}
	if asJSON {
		return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
	}
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "--- %s answered (status %d)\n%s\n\n", v.Original.ServedModel, v.Original.Status,
		or(v.Original.Response, "(no text)"))
	fmt.Fprintf(out, "--- %s answers (status %d, %d ms)\n%s\n", v.Model, v.Status, v.TookMs, or(replyText(v.Body), string(v.Body)))
	return nil
}

// replyText is a completion's text, or "" when it is not one.
func replyText(body json.RawMessage) string {
	var c struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
			Text string `json:"text"`
		} `json:"choices"`
	}
	if json.Unmarshal(body, &c) != nil || len(c.Choices) == 0 {
		return ""
	}
	return c.Choices[0].Message.Content + c.Choices[0].Text
}

func indentJSON(s string) string {
	var v any
	if json.Unmarshal([]byte(s), &v) != nil {
		return s
	}
	b, _ := json.MarshalIndent(v, "", "  ")
	return string(b)
}