code.oscarkilo.com/okg/embed.go

package main

import "flag"
import "fmt"
import "io"
import "os"

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

// runEmbed converts stdin text into embedding vectors, written to
// stdout one vector per line, space-separated floats. Mirrors the
// (now-deprecated) `klex-git/embed` binary.
func runEmbed(cfg *Config, args []string) error {
  fs := flag.NewFlagSet("embed", flag.ContinueOnError)
  model := fs.String("model",
    "openai:text-embedding-3-small",
    "embedding model name")
  dims := fs.Int("dims", 1536,
    "number of vector dimensions to return")
  fullPath := fs.Bool("full-path", false,
    "return one vector per prefix of input "+
      "(otherwise just the final vector)")
  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)

  text, err := io.ReadAll(os.Stdin)
  if err != nil {
    return fmt.Errorf("read stdin: %v", err)
  }

  vectors, err := client.Embed(api.EmbedRequest{
    Text:     string(text),
    Model:    *model,
    Dims:     *dims,
    FullPath: *fullPath,
  })
  if err != nil {
    return fmt.Errorf("embed: %v", err)
  }

  for _, vector := range vectors {
    for i, w := range vector {
      if i > 0 {
        fmt.Print(" ")
      }
      fmt.Printf("%g", w)
    }
    fmt.Println()
  }
  return nil
}