code.oscarkilo.com/klex-git/one/main.go

package main

// This binary runs one LLM inference on one input.

import "encoding/base64"
import "flag"
import "fmt"
import "encoding/json"
import "io/ioutil"
import "log"
import "net/http"
import "os"
import "strings"

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

var model = flag.String("model", "", "overrides .Model, if non-empty")
var system = flag.String("system_file", "", "overrides .System, if non-empty")
var prompt = flag.String("prompt_file", "", "appends to .Messages")
var image = flag.String("image_file", "", "attaches an image to the prompt")
var format = flag.String("format", "text", "text|json|jsonindent")

// guessMimeType returns the MIME type inferred from file contents.
func guessMimeType(b []byte) string {
  if len(b) > 512 {
    b = b[:512]
  }
  return http.DetectContentType(b)
}

func main() {
  flag.Parse()

  // Find the API keys and configure a Klex client.
  config, err := config.ReadConfig()
  if err != nil {
    log.Fatalf("Failed to read config: %v", err)
  }
  client := api.NewClient(config.KlexUrl, config.ApiKey)
  if client == nil {
    log.Fatalf("Failed to create Klex client")
  }

  // Parse stdin as a MessagesRequest object, allowing empty input.
  sin, err := ioutil.ReadAll(os.Stdin)
  if err != nil {
    log.Fatalf("Failed to read stdin: %v", err)
  }
  if len(sin) == 0 {
    sin = []byte("{}")
  }
  var req api.MessagesRequest
  err = json.Unmarshal(sin, &req)
  if err != nil {
    log.Fatalf("Failed to parse a MessagesRequest from stdin: %v", err)
  }

  // Use flags to override parts of the request.
  if *model != "" {
    req.Model = *model
  }
  if *system != "" {
    s, err := ioutil.ReadFile(*system)
    if err != nil {
      log.Fatalf("Failed to read --system_file %s: %v", *system, err)
    }
    req.System = string(s)
  }
  if *image != "" && *prompt == "" {
    log.Fatalf("--image_file requires a non-empty --prompt_file, too")
  }
  if *prompt != "" {
    msg := api.ChatMessage{Role: "user"}
    if *image != "" {
      i, err := ioutil.ReadFile(*image)
      if err != nil {
        log.Fatalf("Failed to read --image_file %s: %v", *image, err)
      }
      mime_type := guessMimeType(i)
      switch mime_type {
      case "image/jpeg", "image/png", "image/gif", "image/webp":
      default:
        log.Fatalf("Unsupported image type: %s", mime_type)
      }
      msg.Content = append(msg.Content, api.ContentBlock{
        Type: "image",
        Source: &api.ContentSource{
          Type: "base64",
          MediaType: mime_type,
          Data: base64.StdEncoding.EncodeToString(i),
        },
      })
    }
    p, err := ioutil.ReadFile(*prompt)
    if err != nil {
      log.Fatalf("Failed to read --prompt_file %s: %v", *prompt, err)
    }
    msg.Content = append(msg.Content, api.ContentBlock{
      Type: "text",
      Text: string(p),
    })
    req.Messages = append(req.Messages, msg)
  }

  // Get LLM output from Klex.
  res, err := client.Messages(req)
  if err != nil {
    log.Fatalf("Klex f() failure: %v", err)
  }

  // Print according to the --format flag.
  out, err := formatResponse(res)
  if err != nil {
    log.Fatalf("Failed to format response: %v", err)
  }
  fmt.Print(out)
}

func formatResponse(res *api.MessagesResponse) (string, error) {
  switch *format {
  case "text":
    var content []string
    for _, c := range res.Content {
      if c.Type == "text" {
        content = append(content, c.Text + "\n")
      }
    }
    return strings.Join(content, "\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)
  }
}