homelab/review-bot/mistral.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Mistral runs the review as a Mistral agent (one per model) with the read-only Gitea connector.
// The conversation is streamed: a request that waits silently for the whole review is cut off
// by Mistral's edge after ten minutes.
type Mistral struct {
	base        string
	key         string
	agentPrefix string
	connectorID string
	http        *http.Client
	// recoverTries and recoverWait bound how often the history of a broken-off conversation is
	// read to find its answer.
	recoverTries int
	recoverWait  time.Duration

	mu     sync.RWMutex
	agents map[string]string // model -> agent id
}

type Agent struct {
	ID          string `json:"id"`
	Name        string `json:"name"`
	Model       string `json:"model"`
	Description string `json:"description"`
}

func (m *Mistral) request(ctx context.Context, method, path string, in any) (*http.Request, error) {
	var body io.Reader
	if in != nil {
		b, err := json.Marshal(in)
		if err != nil {
			return nil, err
		}
		body = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, method, m.base+path, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+m.key)
	if in != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	return req, nil
}

func (m *Mistral) do(ctx context.Context, method, path string, in, out any) error {
	req, err := m.request(ctx, method, path, in)
	if err != nil {
		return err
	}
	resp, err := m.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20))
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusOK {
		return newAPIError("mistral", method, path, resp, data)
	}
	if out == nil {
		return nil
	}
	return json.Unmarshal(data, out)
}

func promptHash(prompt string) string {
	sum := sha256.Sum256([]byte(prompt))
	return hex.EncodeToString(sum[:])[:12]
}

func (m *Mistral) agentName(model string) string { return m.agentPrefix + "-" + model }

// Prepare makes one agent per model match this binary.
func (m *Mistral) Prepare(ctx context.Context, models []string) error {
	for _, model := range models {
		a, err := m.EnsureAgent(ctx, m.agentName(model), model, m.connectorID, reviewPrompt)
		if err != nil {
			return fmt.Errorf("agent for %s: %w", model, err)
		}
		m.mu.Lock()
		if m.agents == nil {
			m.agents = map[string]string{}
		}
		m.agents[model] = a.ID
		m.mu.Unlock()
	}
	return nil
}

func (m *Mistral) agent(model string) string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.agents[model]
}

// EnsureAgent makes the listed agent match this binary: it creates the agent when it does not
// exist and updates it when its model or the hash of its prompt differs, so git stays the source
// of truth for the prompt.
func (m *Mistral) EnsureAgent(ctx context.Context, name, model, connectorID, prompt string) (Agent, error) {
	desc := "PR review, read-only Gitea connector. prompt sha256:" + promptHash(prompt)
	spec := map[string]any{
		"name":         name,
		"model":        model,
		"description":  desc,
		"instructions": prompt,
		"tools":        []map[string]string{{"type": "connector", "connector_id": connectorID}},
		"completion_args": map[string]any{
			"temperature": 0.1,
			"max_tokens":  64000,
		},
	}

	var agents []Agent
	if err := m.do(ctx, http.MethodGet, "/agents?page_size=100", nil, &agents); err != nil {
		return Agent{}, err
	}
	for _, a := range agents {
		if a.Name != name {
			continue
		}
		if a.Model == model && a.Description == desc {
			return a, nil
		}
		var updated Agent
		if err := m.do(ctx, http.MethodPatch, "/agents/"+a.ID, spec, &updated); err != nil {
			return Agent{}, err
		}
		updated.ID = a.ID
		return updated, nil
	}
	var created Agent
	if err := m.do(ctx, http.MethodPost, "/agents", spec, &created); err != nil {
		return Agent{}, err
	}
	return created, nil
}

// A tool is read-only when its name contains _get_, _list_ or _search_, or ends in _read.
var readOnlyTool = regexp.MustCompile(`_(get|list|search)_|_read$`)

// answerBuilder collects the answer from conversation outputs: the text of message outputs
// (thinking is dropped), and the tool calls, listing every one that is not read-only.
type answerBuilder struct {
	a      Answer
	texts  []string
	lastID string
}

func (b *answerBuilder) tool(name string) {
	b.a.ToolCalls++
	if !readOnlyTool.MatchString(name) {
		b.a.Violations = append(b.a.Violations, name)
	}
}

// text adds content (a string, a chunk or a list of chunks) of message id. Content of the same
// message is joined directly, different messages are joined by a newline.
func (b *answerBuilder) text(id string, content json.RawMessage) {
	var parts []string
	var s string
	if json.Unmarshal(content, &s) == nil {
		parts = append(parts, s)
	} else {
		type chunk struct {
			Type string `json:"type"`
			Text string `json:"text"`
		}
		var chunks []chunk
		if json.Unmarshal(content, &chunks) != nil {
			var one chunk
			if json.Unmarshal(content, &one) == nil {
				chunks = []chunk{one}
			}
		}
		for _, ch := range chunks {
			if ch.Type == "text" {
				parts = append(parts, ch.Text)
			}
		}
	}
	if len(parts) == 0 {
		return
	}
	joined := strings.Join(parts, "")
	if id != "" && id == b.lastID && len(b.texts) > 0 {
		b.texts[len(b.texts)-1] += joined
		return
	}
	b.texts = append(b.texts, joined)
	b.lastID = id
}

func (b *answerBuilder) answer() Answer {
	b.a.Text = strings.Join(b.texts, "\n")
	return b.a
}

type entry struct {
	Type    string          `json:"type"`
	ID      string          `json:"id"`
	Name    string          `json:"name"`
	Content json.RawMessage `json:"content"`
}

// parseAnswer reads the answer from a conversation response (`outputs`) or history (`entries`).
func parseAnswer(raw []byte) (Answer, error) {
	var c struct {
		Outputs []entry `json:"outputs"`
		Entries []entry `json:"entries"`
	}
	if err := json.Unmarshal(raw, &c); err != nil {
		return Answer{}, err
	}
	var b answerBuilder
	for _, o := range append(c.Outputs, c.Entries...) {
		switch o.Type {
		case "tool.execution":
			b.tool(o.Name)
		case "message.output":
			// Every entry is its own message.
			b.lastID = ""
			b.text("", o.Content)
		}
	}
	return b.answer(), nil
}

type event struct {
	Type           string          `json:"type"`
	ConversationID string          `json:"conversation_id"`
	ID             string          `json:"id"`
	Name           string          `json:"name"`
	Content        json.RawMessage `json:"content"`
	Message        string          `json:"message"`
	Code           int             `json:"code"`
}

// Ask streams one conversation with the model's agent. When the stream breaks off after the
// conversation started, the answer is looked up in the stored history before giving up.
func (m *Mistral) Ask(ctx context.Context, model, task string) (Answer, error) {
	agentID := m.agent(model)
	if agentID == "" {
		return Answer{}, fmt.Errorf("no Mistral agent for %s yet", model)
	}
	in := map[string]any{"agent_id": agentID, "inputs": task, "store": true, "stream": true}
	req, err := m.request(ctx, http.MethodPost, "/conversations", in)
	if err != nil {
		return Answer{}, err
	}
	req.Header.Set("Accept", "text/event-stream")
	resp, err := m.http.Do(req)
	if err != nil {
		return Answer{}, &streamError{Provider: "mistral", Err: err}
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		data, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
		return Answer{}, newAPIError("mistral", http.MethodPost, "/conversations", resp, data)
	}
	conv, ans, err := readStream(resp.Body)
	if err == nil {
		return ans, nil
	}
	var broken *streamError
	if conv == "" || !errors.As(err, &broken) || ctx.Err() != nil {
		return Answer{}, err
	}
	if rec, ok := m.recover(ctx, conv); ok {
		return rec, nil
	}
	return Answer{}, err
}

// readStream reads the server-sent events of a conversation until conversation.response.done.
// It returns the conversation id as soon as it is known, also with an error.
func readStream(body io.Reader) (conv string, _ Answer, _ error) {
	r := bufio.NewReaderSize(body, 64<<10)
	var b answerBuilder
	var data []string
	dispatch := func() (done bool, err error) {
		if len(data) == 0 {
			return false, nil
		}
		raw := strings.Join(data, "\n")
		data = data[:0]
		var ev event
		if err := json.Unmarshal([]byte(raw), &ev); err != nil {
			return false, &streamError{Provider: "mistral", Err: fmt.Errorf("unreadable event: %w", err)}
		}
		switch ev.Type {
		case "conversation.response.started":
			conv = ev.ConversationID
		case "message.output.delta":
			b.text(ev.ID, ev.Content)
		case "tool.execution.started":
			b.tool(ev.Name)
		case "conversation.response.error":
			return false, fmt.Errorf("mistral conversation error %d: %s", ev.Code, ev.Message)
		case "conversation.response.done":
			return true, nil
		}
		return false, nil
	}
	for {
		line, err := r.ReadString('\n')
		line = strings.TrimRight(line, "\r\n")
		switch {
		case line == "":
			if done, derr := dispatch(); derr != nil {
				return conv, Answer{}, derr
			} else if done {
				return conv, b.answer(), nil
			}
		case strings.HasPrefix(line, "data:"):
			data = append(data, strings.TrimPrefix(strings.TrimPrefix(line, "data:"), " "))
		}
		if err != nil {
			if done, derr := dispatch(); derr != nil {
				return conv, Answer{}, derr
			} else if done {
				return conv, b.answer(), nil
			}
			if err == io.EOF {
				err = io.ErrUnexpectedEOF
			}
			return conv, Answer{}, &streamError{Provider: "mistral", Err: err}
		}
	}
}

// recover reads the stored history of a conversation whose stream broke off, a few times, and
// takes its answer once it ends in a verdict.
func (m *Mistral) recover(ctx context.Context, conv string) (Answer, bool) {
	for try := 0; try < m.recoverTries; try++ {
		select {
		case <-ctx.Done():
			return Answer{}, false
		case <-time.After(m.recoverWait):
		}
		var raw json.RawMessage
		if err := m.do(ctx, http.MethodGet, "/conversations/"+conv+"/history", nil, &raw); err != nil {
			continue
		}
		if a, err := parseAnswer(raw); err == nil {
			if _, verdict := splitVerdict(a.Text); verdict != "" {
				return a, true
			}
		}
	}
	return Answer{}, false
}