llmgw/internal/cli/answercache.go

package cli

import (
	"fmt"
	"time"

	"github.com/spf13/cobra"

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

func newAnswerCacheCmd(open opener) *cobra.Command {
	var (
		project   string
		ttl       time.Duration
		anyAnswer bool
		similar   float64
		off       bool
	)
	cmd := &cobra.Command{
		Use:   "cache",
		Short: "Reuse a project's answers: the same request is answered again without the model",
		Long: "With this on, a request of the project that is the same as one answered within\n" +
			"--ttl, to the same model, gets that answer again, from the gateway's memory,\n" +
			"rather than asking the model. It is answered with x-llmgw-cache: hit, recorded\n" +
			"as a request with no tokens, and costs nothing against a budget.\n\n" +
			"Only answers that would come out the same are reused: a chat or completion\n" +
			"asked at temperature 0 or with a seed, and embeddings, rerank and score.\n" +
			"--any reuses every answer, sampled or not. A client sending\n" +
			"Cache-Control: no-cache or x-llmgw-cache: no is always asked afresh, and so is\n" +
			"a request continuing a stored conversation or under a guardrails no_store\n" +
			"policy. The guardrails judge a reused answer as they judge any other.\n\n" +
			"Off for every project until this is run. Without --ttl or --off, says how\n" +
			"it stands. The answers are held in memory, router.cache bounds how much.\n\n" +
			"--similar also answers a chat question that is like one already answered:\n" +
			"its last user message embedded by router.cache.embedding_model, and given\n" +
			"the answer whose question reaches that cosine similarity (0.95 is strict,\n" +
			"0.9 looser), as x-llmgw-cache: similar. The system prompt, the earlier\n" +
			"turns, tools and the sampling settings must still be the same.",
		Example: "  llmgw keys cache --project Acme/ml/train --ttl 1h\n" +
			"  llmgw keys cache --project Acme/ml/eval --ttl 24h --any\n" +
			"  llmgw keys cache --project Acme/support/desk --ttl 24h --similar 0.95\n" +
			"  llmgw keys cache --project Acme/ml/train --off",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			switch {
			case off && (ttl != 0 || anyAnswer || similar != 0):
				return fmt.Errorf("--off, or --ttl with --any or --similar: not both")
			case similar < 0 || similar > 1:
				return fmt.Errorf("--similar is a cosine similarity between 0 and 1, like 0.95")
			case similar != 0 && ttl == 0:
				return fmt.Errorf("--similar goes with --ttl: say how long answers are reused")
			case ttl < 0:
				return fmt.Errorf("--ttl is how long an answer is reused, like 1h")
			case ttl > 0 && ttl < time.Second:
				return fmt.Errorf("--ttl is at least a second")
			case anyAnswer && ttl == 0:
				return fmt.Errorf("--any goes with --ttl: say how long answers are reused")
			}
			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 || ttl > 0 {
				if err := s.SetAnswerCache(cmd.Context(), p.ID, store.AnswerCache{TTL: ttl, Any: anyAnswer, Similar: similar}); err != nil {
					return err
				}
			}
			c, err := s.AnswerCacheOf(cmd.Context(), p.ID)
			if err != nil {
				return err
			}
			switch {
			case c.TTL == 0:
				fmt.Fprintf(cmd.OutOrStdout(), "%s reuses no answers.\n", p.Path)
			case c.Any:
				fmt.Fprintf(cmd.OutOrStdout(), "%s reuses every answer for %s, sampled or not.\n", p.Path, c.TTL)
			default:
				fmt.Fprintf(cmd.OutOrStdout(), "%s reuses answers for %s: those asked at temperature 0 or with a seed, "+
					"and embeddings, rerank and score.\n", p.Path, c.TTL)
			}
			if c.TTL > 0 && c.Similar > 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "Chat questions reaching %.2f similarity to one answered get its answer too.\n", c.Similar)
			}
			if off || ttl > 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "Gateways see the change within 30 seconds.")
			}
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "", "project: id, path or unique name")
	cmd.Flags().DurationVar(&ttl, "ttl", 0, "how long an answer is reused, like 1h")
	cmd.Flags().BoolVar(&anyAnswer, "any", false, "reuse sampled answers too, not only reproducible ones")
	cmd.Flags().Float64Var(&similar, "similar", 0, "also answer chat questions this similar to one answered, like 0.95")
	cmd.Flags().BoolVar(&off, "off", false, "stop reusing answers for this project")
	return cmd
}