code.oscarkilo.com/okg/one.go

package main

import "encoding/json"
import "flag"
import "fmt"
import "io"
import "os"
import "strings"

import "oscarkilo.com/klex-git/api"

// runOne runs one LLM inference on one input. Mirrors the (now-
// deprecated) `klex-git/one` binary.
//
// Reads stdin as an api.MessagesRequest JSON; empty stdin is
// allowed (treated as {}). Flags override individual fields.
func runOne(cfg *Config, args []string) error {
  fs := flag.NewFlagSet("one", flag.ContinueOnError)
  model := fs.String("model", "",
    "override .Model, if non-empty")
  systemFile := fs.String("system-file", "",
    "override .System with the contents of this file")
  promptFile := fs.String("prompt-file", "",
    "append this file to .Messages as a user prompt")
  attach := fs.String("attach", "",
    "path to a file (image or PDF) to attach to the prompt")
  format := fs.String("format", "text",
    "text | json | jsonindent")
  fastFail := fs.Bool("fast-fail", true,
    "preflight-check the attachment MIME against the model's "+
      "llm2 capabilities; fail before paying for the call. "+
      "Set false in tight loops to skip the extra HTTP "+
      "round-trip per call.")
  if err := fs.Parse(args); err != nil {
    return err
  }

  if cfg.ApiKey == "" {
    return fmt.Errorf(
      "no API key — run `okg auth login --key sk-...`")
  }
  client := newKlexClient(cfg)

  // Parse stdin as a MessagesRequest; empty → {}.
  sin, err := io.ReadAll(os.Stdin)
  if err != nil {
    return fmt.Errorf("read stdin: %v", err)
  }
  if len(sin) == 0 {
    sin = []byte("{}")
  }
  var req api.MessagesRequest
  if err := json.Unmarshal(sin, &req); err != nil {
    return fmt.Errorf("parse MessagesRequest: %v", err)
  }

  // Flag overrides.
  if *model != "" {
    req.Model = *model
  }
  if *systemFile != "" {
    s, err := os.ReadFile(*systemFile)
    if err != nil {
      return fmt.Errorf(
        "read --system-file %s: %v", *systemFile, err)
    }
    req.System = string(s)
  }
  if *attach != "" && *promptFile == "" {
    return fmt.Errorf(
      "--attach requires a non-empty --prompt-file")
  }
  var attachMime string
  if *promptFile != "" {
    msg := api.ChatMessage{Role: "user"}
    if *attach != "" {
      data, err := os.ReadFile(*attach)
      if err != nil {
        return fmt.Errorf(
          "read --attach %s: %v", *attach, err)
      }
      blk := api.NewDocumentBlock(data)
      attachMime = blk.Source.MediaType
      msg.Content = append(msg.Content, blk)
    }
    p, err := os.ReadFile(*promptFile)
    if err != nil {
      return fmt.Errorf(
        "read --prompt-file %s: %v", *promptFile, err)
    }
    msg.Content = append(msg.Content, api.ContentBlock{
      Type: "text",
      Text: string(p),
    })
    req.Messages = append(req.Messages, msg)
  }

  // Preflight: catch unsupported attachment types before the
  // call.
  if *fastFail && attachMime != "" {
    if err := preflightAttachment(
      client, req.Model, attachMime,
    ); err != nil {
      return err
    }
  }

  res, err := client.Messages(req)
  if err != nil {
    return fmt.Errorf("klex f() failed: %v", err)
  }

  out, err := formatMessagesResponse(res, *format)
  if err != nil {
    return err
  }
  fmt.Print(out)
  return nil
}

func formatMessagesResponse(
  res *api.MessagesResponse, format string,
) (string, error) {
  switch format {
  case "text":
    var parts []string
    for _, c := range res.Content {
      if c.Type == "text" {
        parts = append(parts, c.Text+"\n")
      }
    }
    return strings.Join(parts, "\n"), nil
  case "json":
    buf, err := json.Marshal(res)
    return string(buf), err
  case "jsonindent":
    buf, err := json.MarshalIndent(res, "", "  ")
    return string(buf), err
  default:
    return "", fmt.Errorf(
      "unsupported --format=%s", format)
  }
}

// preflightAttachment fetches the model's llm2 capabilities and
// returns an error if it can't accept the given attachment MIME
// type. Returns nil silently for MIME families Klex has no
// capability flag for (anything that isn't image/* or
// application/pdf).
func preflightAttachment(
  client *api.Client, modelName, mimeType string,
) error {
  resp, err := client.ListFuncs("latest")
  if err != nil {
    return fmt.Errorf(
      "preflight ListFuncs failed "+
        "(--fast-fail=false to bypass): %v", err)
  }
  var fn *api.Func
  for i := range resp.Funcs {
    if resp.Funcs[i].Name == modelName {
      fn = &resp.Funcs[i]
      break
    }
  }
  if fn == nil {
    return fmt.Errorf(
      "unknown model %q (--fast-fail=false to bypass)",
      modelName)
  }
  if len(fn.Versions) == 0 ||
    fn.Versions[len(fn.Versions)-1].LLM2 == nil {
    return fmt.Errorf(
      "model %q has no llm2 config", modelName)
  }
  llm := fn.Versions[len(fn.Versions)-1].LLM2
  switch {
  case strings.HasPrefix(mimeType, "image/"):
    if !llm.CanSeeImages {
      return fmt.Errorf(
        "model %q does not accept images "+
          "(can_see_images=false)", modelName)
    }
  case mimeType == "application/pdf":
    if !llm.CanSeePDFs {
      return fmt.Errorf(
        "model %q does not accept PDFs "+
          "(can_see_pdfs=false)", modelName)
    }
  }
  return nil
}