llmgw/internal/cli/requestlog.go

package cli

import (
	"fmt"

	"github.com/spf13/cobra"
)

func newRequestLogCmd(open opener) *cobra.Command {
	var (
		project string
		hours   int
		sample  float64
		off     bool
	)
	cmd := &cobra.Command{
		Use:   "requests",
		Short: "Keep a project's requests and answers, to follow one up or replay it",
		Long: "With this on, what each request of the project sent the model — after the\n" +
			"guardrails' redaction, with any stored conversation in it — and the answer\n" +
			"the client got are kept for --retain-hours. `llmgw ctl request ID --body`\n" +
			"shows one, and `--replay MODEL` sends it to another model.\n\n" +
			"Off for every project until this is run, like conversations: llmgw\n" +
			"otherwise keeps no text. --sample keeps a share of the requests rather\n" +
			"than all of them. A project under a guardrails policy with no_store keeps\n" +
			"nothing whatever this says.",
		Example: "  llmgw keys requests --project Acme/ml/train --retain-hours 72\n" +
			"  llmgw keys requests --project Acme/ml/train --retain-hours 168 --sample 0.05\n" +
			"  llmgw keys requests --project Acme/ml/train --off",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if !off && hours <= 0 {
				return fmt.Errorf("--retain-hours is required: say how long requests are kept, or pass --off")
			}
			if sample <= 0 || sample > 1 {
				return fmt.Errorf("--sample is a share: above 0, at most 1")
			}
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()
			p, err := resolveProject(cmd.Context(), s, project)
			if err != nil {
				return err
			}
			if off {
				hours = 0
			}
			if err := s.SetRequestLog(cmd.Context(), p.ID, hours, sample); err != nil {
				return err
			}
			if hours == 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "%s keeps no requests. Those kept are removed at the next prune.\n", p.Path)
				return nil
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s keeps %s of its requests and answers for %d hours.\n",
				p.Path, share(sample), hours)
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "", "project: id, path or unique name")
	cmd.Flags().IntVar(&hours, "retain-hours", 0, "how long a request and its answer are kept")
	cmd.Flags().Float64Var(&sample, "sample", 1, "the share of requests kept, above 0 and at most 1")
	cmd.Flags().BoolVar(&off, "off", false, "stop keeping requests for this project")
	return cmd
}

func share(f float64) string {
	if f >= 1 {
		return "every one"
	}
	return fmt.Sprintf("%g%%", f*100)
}

func newResidencyCmd(open opener) *cobra.Command {
	var (
		project           string
		inCluster, anyway bool
	)
	cmd := &cobra.Command{
		Use:   "residency",
		Short: "Keep a project's requests in the cluster, never sent to an outside provider",
		Long: "With --in-cluster, a model served by an outside provider (external:) is\n" +
			"refused to the project's keys, and never used for them as a fallback or\n" +
			"a canary: its prompts do not leave the machines the cluster runs on.\n" +
			"--anywhere lifts it. Without either, says how it stands.",
		Example: "  llmgw keys residency --project Acme/legal --in-cluster",
		Args:    cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if inCluster && anyway {
				return fmt.Errorf("--in-cluster or --anywhere, not both")
			}
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()
			p, err := resolveProject(cmd.Context(), s, project)
			if err != nil {
				return err
			}
			if inCluster || anyway {
				if err := s.SetInCluster(cmd.Context(), p.ID, inCluster); err != nil {
					return err
				}
			}
			on, err := s.InCluster(cmd.Context(), p.ID)
			if err != nil {
				return err
			}
			if on {
				fmt.Fprintf(cmd.OutOrStdout(), "%s's requests stay in the cluster.\n", p.Path)
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s's requests may go to outside providers its keys allow.\n", p.Path)
			}
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "", "project: id, path or unique name")
	cmd.Flags().BoolVar(&inCluster, "in-cluster", false, "keep its requests in the cluster")
	cmd.Flags().BoolVar(&anyway, "anywhere", false, "let them go to outside providers again")
	return cmd
}