homelab/review-bot/selector.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
)

// PRFile is one changed file of a pull request.
type PRFile struct {
	Filename  string `json:"filename"`
	Status    string `json:"status"`
	Additions int    `json:"additions"`
	Deletions int    `json:"deletions"`
}

// Choice is the selector's pick with the probability of every model.
type Choice struct {
	Model         string
	Confidence    float64
	Probabilities map[string]float64
}

// selectInstructions is the question the selector answers for every pull request.
const selectInstructions = "Pick the model that should review the pull request in `pull_request`. " +
	"The review reads the repository, checks the change for bugs, security problems and broken behaviour, " +
	"and must be thorough. Prefer a cheaper or faster model when the change is small and low-risk " +
	"(documentation, version bumps, simple configuration). Pick a stronger model when the change is large, " +
	"spans several components, or touches security, authentication, secrets, network policy, concurrency, " +
	"data migrations or anything that is hard to undo."

// TypeSafe picks the review model with a System One Choice question.
type TypeSafe struct {
	url   string
	key   string
	model string
	http  *http.Client
}

func (t *TypeSafe) Select(ctx context.Context, slug string, pr PR, files []PRFile, models []Model) (Choice, error) {
	criteria := map[string]any{}
	for _, m := range models {
		if m.Description != "" {
			criteria[m.ID()] = m.Description
		} else {
			criteria[m.ID()] = nil
		}
	}
	state := map[string]any{"pull_request": map[string]any{
		"repository": slug, "title": pr.Title, "description": pr.Body,
		"additions": pr.Additions, "deletions": pr.Deletions, "changed_files": pr.ChangedFiles,
		"files": files,
	}}
	in := map[string]any{
		"state": state,
		"model": t.model,
		"questions": map[string]any{"model": map[string]any{
			"type": "choice", "instructions": selectInstructions, "criteria": criteria,
		}},
	}
	b, err := json.Marshal(in)
	if err != nil {
		return Choice{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, t.url, bytes.NewReader(b))
	if err != nil {
		return Choice{}, err
	}
	req.Header.Set("Authorization", "Bearer "+t.key)
	req.Header.Set("Content-Type", "application/json")
	resp, err := t.http.Do(req)
	if err != nil {
		return Choice{}, err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return Choice{}, err
	}
	if resp.StatusCode != http.StatusOK {
		return Choice{}, newAPIError("typesafe", http.MethodPost, "/v1/systemone", resp, data)
	}
	var out struct {
		Answers map[string]struct {
			Choice        string             `json:"choice"`
			Confidence    float64            `json:"confidence"`
			Probabilities map[string]float64 `json:"probabilities"`
		} `json:"answers"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return Choice{}, err
	}
	a, ok := out.Answers["model"]
	if !ok || a.Choice == "" {
		return Choice{}, fmt.Errorf("typesafe answered without a model: %s", snippet(data))
	}
	return Choice{Model: a.Choice, Confidence: a.Confidence, Probabilities: a.Probabilities}, nil
}