so-many-crows/jev.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"
)

// Jev (TypeSafe) is a classifier, not a text model: it returns typed choices
// with calibrated confidence. The Sentry uses it for triage, and escalates to
// the LLM when Jev is unsure, when an item needs subtasks written, or when Jev
// fails. Docs: https://docs.typesafe.ai/introduction/quickstart

type JevClient struct {
	url, key, model string
	threshold       float64 // both answers must reach this confidence to stand alone
	client          *http.Client
}

// JevCall is one classification: a label and a confidence per question.
type JevCall struct {
	Mode, TaskType         string
	ModeConf, TaskTypeConf float64
	Model                  string
	Tokens                 int
}

// NewJev returns nil (no classifier) unless TYPESAFE_API_KEY is set.
//
//	TYPESAFE_API_KEY         the key from the TypeSafe console
//	CROWS_JEV_MODEL          default jev-latest
//	CROWS_JEV_THRESHOLD      default 0.7
func NewJev() *JevClient {
	key := os.Getenv("TYPESAFE_API_KEY")
	if key == "" {
		return nil
	}
	model := os.Getenv("CROWS_JEV_MODEL")
	if model == "" {
		model = "jev-latest"
	}
	threshold := 0.7
	if v, err := strconv.ParseFloat(os.Getenv("CROWS_JEV_THRESHOLD"), 64); err == nil && v > 0 && v <= 1 {
		threshold = v
	}
	return &JevClient{url: "https://api.typesafe.ai/v1/systemone", key: key, model: model, threshold: threshold, client: &http.Client{}}
}

func (j *JevClient) Name() string {
	return fmt.Sprintf("Jev (%s, escalates below %.0f%%)", j.model, j.threshold*100)
}

var jevQuestions = map[string]any{
	"mode": map[string]any{
		"type":         "choice",
		"instructions": "How should a team of AI agents handle this todo item?",
		"criteria": map[string]string{
			"solo":   "Simple or single-step: one agent can answer it directly.",
			"split":  "Large, with 2 to 4 distinct parts that can be done independently and in parallel.",
			"murder": "Ambiguous, high-stakes or a judgment call: several agents should answer independently, then reconcile.",
		},
	},
	"task_type": map[string]any{
		"type":         "choice",
		"instructions": "What kind of task is this todo item?",
		"criteria": map[string]string{
			"retrieval": "Looking up or recalling facts.",
			"writing":   "Drafting or editing text.",
			"analysis":  "Comparing, evaluating, deciding or reasoning about options.",
			"other":     "Anything else.",
		},
	},
}

func (j *JevClient) Classify(ctx context.Context, item string) (JevCall, error) {
	body, _ := json.Marshal(map[string]any{"model": j.model, "state": item, "questions": jevQuestions})
	req, err := http.NewRequestWithContext(ctx, "POST", j.url, bytes.NewReader(body))
	if err != nil {
		return JevCall{}, err
	}
	req.Header.Set("Authorization", "Bearer "+j.key)
	req.Header.Set("Content-Type", "application/json")
	resp, err := j.client.Do(req)
	if err != nil {
		return JevCall{}, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return JevCall{}, fmt.Errorf("Jev returned %d: %s", resp.StatusCode, clip(raw, 300))
	}
	type answer struct {
		Choice     string  `json:"choice"`
		Confidence float64 `json:"confidence"`
	}
	var out struct {
		Model   string            `json:"model"`
		Answers map[string]answer `json:"answers"`
		Usage   struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return JevCall{}, fmt.Errorf("Jev sent an unreadable reply: %w", err)
	}
	mode, task := out.Answers["mode"], out.Answers["task_type"]
	if mode.Choice == "" || task.Choice == "" {
		return JevCall{}, errors.New("Jev's reply was missing an answer")
	}
	model := out.Model
	if model == "" {
		model = j.model
	}
	return JevCall{
		Mode: mode.Choice, TaskType: task.Choice, ModeConf: mode.Confidence, TaskTypeConf: task.Confidence,
		Model: model, Tokens: out.Usage.InputTokens + out.Usage.OutputTokens,
	}, nil
}

func (c JevCall) confident(threshold float64) bool {
	return c.ModeConf >= threshold && c.TaskTypeConf >= threshold
}

func (c JevCall) summary() string {
	return fmt.Sprintf("Jev: %s %.0f%%, %s %.0f%%", c.Mode, c.ModeConf*100, c.TaskType, c.TaskTypeConf*100)
}

// classify runs Jev under the same guardrails as any crow: budget, flock cap,
// timeout, flight log, and flying/landed events for the Sentry.
func (f *Flock) classify(ctx context.Context, w Watcher, item string, budgets ...*Budget) (JevCall, error) {
	for _, b := range budgets {
		if err := b.Check(); err != nil {
			return JevCall{}, err
		}
	}
	f.sem <- struct{}{}
	defer func() { <-f.sem }()
	if w != nil {
		w.Flying("Sentry")
	}
	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	call, err := f.jev.Classify(cctx, item)
	for _, b := range budgets {
		b.Add(call.Tokens)
	}
	e := Event{Kind: "flight", Crow: "Sentry", Model: call.Model, Tokens: call.Tokens, Ms: time.Since(start).Milliseconds()}
	if err != nil {
		e.Error = err.Error()
	} else {
		e.Note = call.summary()
	}
	if f.log != nil {
		f.log.Write(e)
	}
	if w != nil {
		w.Landed("Sentry", call.Tokens, err)
	}
	return call, err
}