email_agent/internal/classify/jev.go

// Package classify turns stored messages into a category, an importance, a
// suggested action, and a confidence using the TypeSafe (JEV) HTTP API, then
// applies the user's rules (FR-3, FR-4).
//
// The JEV request/response shape is confirmed against the live docs
// (https://docs.typesafe.ai/api.md) and recorded in docs/phase3-classify.md:
//
//	POST https://api.typesafe.ai/v1/systemone
//	Authorization: Bearer <API_KEY>
//	{"state": …, "model": "jev-latest", "questions": {"‹id›": {"type": "choice", …}}}
//	→ {"model": "jev-1.13.0", "answers": {"‹id›": {"type": "choice", "confidence", "probabilities"}}}
package classify

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"math/rand"
	"net/http"
	"strings"
	"time"
)

// DefaultBaseURL is TypeSafe's documented API host. JEV_BASE_URL overrides it.
const DefaultBaseURL = "https://api.typesafe.ai"

// DefaultModel is the flagship System One alias documented at
// https://docs.typesafe.ai/models.
const DefaultModel = "jev-latest"

// systemOnePath is the documented evaluation endpoint.
const systemOnePath = "/v1/systemone"

// ErrNotConfigured means JEV_API_KEY is not set.
var ErrNotConfigured = errors.New("JEV_API_KEY is not configured")

// maxBodyChars bounds how much of a message body is sent to JEV.
const maxBodyChars = 4000

// MessageInput is the message text handed to JEV.
type MessageInput struct {
	From     string
	To       string
	Subject  string
	Source   string
	Body     string
	Received string
}

// CategoryOption is one Choice option: the live category name and its rubric.
type CategoryOption struct {
	Name        string
	Description string
}

// Result is one classification decision. Confidence is the Choice answer's own
// confidence, so the UI can show how sure JEV is.
type Result struct {
	Category             string
	CategoryConfidence   float64
	Importance           string
	ImportanceConfidence float64
	Action               string
	ActionConfidence     float64
	Model                string
}

// Client calls the TypeSafe System One endpoint.
type Client struct {
	baseURL string
	apiKey  string
	model   string
	http    *http.Client
	logger  *slog.Logger
}

// NewClient builds a JEV client. An empty baseURL falls back to DefaultBaseURL,
// an empty model to DefaultModel. logger may be nil.
func NewClient(baseURL, apiKey, model string, logger *slog.Logger) *Client {
	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if baseURL == "" {
		baseURL = DefaultBaseURL
	}
	if strings.TrimSpace(model) == "" {
		model = DefaultModel
	}
	if logger == nil {
		logger = slog.New(slog.DiscardHandler)
	}
	return &Client{
		baseURL: baseURL,
		apiKey:  strings.TrimSpace(apiKey),
		model:   model,
		http:    &http.Client{Timeout: 60 * time.Second},
		logger:  logger,
	}
}

// Configured reports whether an API key is present.
func (c *Client) Configured() bool { return c.apiKey != "" }

// Model returns the configured model alias.
func (c *Client) Model() string { return c.model }

// Classify sends one message and returns the category, importance, action, and
// confidences from one JEV request.
func (c *Client) Classify(ctx context.Context, msg MessageInput, categories []CategoryOption) (Result, error) {
	if !c.Configured() {
		return Result{}, ErrNotConfigured
	}
	if len(categories) == 0 {
		return Result{}, errors.New("no categories configured")
	}

	categoryCriteria := make(map[string]any, len(categories))
	for _, opt := range categories {
		if strings.TrimSpace(opt.Name) == "" {
			continue
		}
		if strings.TrimSpace(opt.Description) == "" {
			categoryCriteria[opt.Name] = nil
		} else {
			categoryCriteria[opt.Name] = opt.Description
		}
	}

	body := map[string]any{
		"state": map[string]any{
			"email": map[string]any{
				"from":           msg.From,
				"to":             msg.To,
				"source_address": msg.Source,
				"source_domain":  domainOf(msg.Source),
				"subject":        msg.Subject,
				"received":       msg.Received,
				"body":           truncate(msg.Body, maxBodyChars),
			},
		},
		"model": c.model,
		"questions": map[string]any{
			"category": map[string]any{
				"type":         "choice",
				"instructions": "Which category best fits this email for a personal family inbox? Pick the single best match.",
				"criteria":     categoryCriteria,
			},
			"importance": map[string]any{
				"type":         "choice",
				"instructions": "How important is this email to the recipient?",
				"criteria": map[string]any{
					"low":    "Can be ignored or removed without loss",
					"medium": "Worth a quick look",
					"high":   "Needs the recipient's attention soon",
				},
			},
			"action": map[string]any{
				"type":         "choice",
				"instructions": "What is the best single recommended action for this email?",
				"criteria": map[string]any{
					"delete":      "Junk or spam; safe to remove",
					"unsubscribe": "Stop future mail from this sender",
					"reply":       "Send a reply to the sender",
					"file":        "Keep for reference; move out of the inbox",
					"keep":        "Leave in the inbox",
				},
			},
		},
	}

	resp, err := c.post(ctx, body)
	if err != nil {
		return Result{}, err
	}

	res := Result{Model: resp.Model}
	if a, ok := resp.Answers["category"]; ok {
		res.Category = a.Choice
		res.CategoryConfidence = a.Confidence
	}
	if a, ok := resp.Answers["importance"]; ok {
		res.Importance = a.Choice
		res.ImportanceConfidence = a.Confidence
	}
	if a, ok := resp.Answers["action"]; ok {
		res.Action = a.Choice
		res.ActionConfidence = a.Confidence
	}
	if res.Model == "" {
		res.Model = c.model
	}
	return res, nil
}

// answer is one System One Choice answer.
type answer struct {
	Type          string             `json:"type"`
	Choice        string             `json:"choice"`
	Confidence    float64            `json:"confidence"`
	Probabilities map[string]float64 `json:"probabilities"`
	Noul          float64            `json:"noul"`
}

type response 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"`
}

// post sends the request with bounded retries for rate limits and outages.
func (c *Client) post(ctx context.Context, body map[string]any) (response, error) {
	payload, err := json.Marshal(body)
	if err != nil {
		return response{}, fmt.Errorf("encode JEV request: %w", err)
	}

	const maxAttempts = 4
	var lastErr error
	for attempt := 1; attempt <= maxAttempts; attempt++ {
		resp, retryable, err := c.do(ctx, payload)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if !retryable || attempt == maxAttempts {
			break
		}
		delay := backoff(attempt)
		c.logger.Warn("jev request failed; retrying", "attempt", attempt, "delay_ms", delay.Milliseconds(), "error", err)
		select {
		case <-ctx.Done():
			return response{}, ctx.Err()
		case <-time.After(delay):
		}
	}
	return response{}, lastErr
}

func (c *Client) do(ctx context.Context, payload []byte) (response, bool, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+systemOnePath, bytes.NewReader(payload))
	if err != nil {
		return response{}, false, fmt.Errorf("build JEV request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	httpResp, err := c.http.Do(req)
	if err != nil {
		return response{}, true, fmt.Errorf("call JEV: %w", err)
	}
	defer func() { _ = httpResp.Body.Close() }()

	raw, err := io.ReadAll(io.LimitReader(httpResp.Body, 1<<20))
	if err != nil {
		return response{}, true, fmt.Errorf("read JEV response: %w", err)
	}

	if httpResp.StatusCode != http.StatusOK {
		retryable := httpResp.StatusCode == http.StatusTooManyRequests ||
			httpResp.StatusCode == http.StatusServiceUnavailable ||
			httpResp.StatusCode == 529 ||
			httpResp.StatusCode >= 500
		return response{}, retryable,
			fmt.Errorf("JEV returned %d: %s", httpResp.StatusCode, strings.TrimSpace(string(raw)))
	}

	var out response
	if err := json.Unmarshal(raw, &out); err != nil {
		return response{}, false, fmt.Errorf("decode JEV response: %w", err)
	}
	if len(out.Answers) == 0 {
		return response{}, false, fmt.Errorf("JEV response contained no answers")
	}
	return out, false, nil
}

// backoff returns the delay before the next attempt, with a little jitter so
// parallel workers do not retry in lock step.
func backoff(attempt int) time.Duration {
	base := time.Duration(1<<(attempt-1)) * time.Second
	if base > 8*time.Second {
		base = 8 * time.Second
	}
	return base + time.Duration(rand.Int63n(int64(base/2+1)))
}

func truncate(s string, max int) string {
	s = strings.TrimSpace(s)
	if len(s) <= max {
		return s
	}
	return s[:max] + "…"
}

func domainOf(addr string) string {
	if i := strings.LastIndex(addr, "@"); i >= 0 && i < len(addr)-1 {
		return addr[i+1:]
	}
	return addr
}