BountyBoard/services/atomizer/llm.go

package main

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

// llmClient calls Anthropic via the OpenAI-compatible chat-completions
// endpoint (default) or the native Messages API (LLM_API_STYLE=anthropic).
// Plain net/http, no SDK (§9.1).
type llmClient struct {
	cfg  config
	log  *slog.Logger
	http *http.Client
}

func newLLMClient(cfg config, log *slog.Logger) *llmClient {
	return &llmClient{cfg: cfg, log: log, http: &http.Client{Timeout: 110 * time.Second}}
}

const atomizeSystem = `You are an expert software project planner. You split one ticket into small, well-scoped, independently deliverable developer tasks.
Respond with ONLY a JSON object, no prose and no markdown fences, exactly matching:
{"tasks":[{"title":"...","description":"...","acceptanceCriteria":["..."],"effortCoefficient":0.25}],"notes":"short advice for the consultant"}
Rules: between MIN and MAX tasks; every effortCoefficient is a number in (0,1]; all effortCoefficients MUST sum to exactly 1.0; titles are concise; descriptions are self-contained instructions; each task has 1-4 testable acceptance criteria.`

const extendSystem = `You are an expert software project planner. You design exactly ONE additional sibling task that extends the given task's functionality per the extension note.
Respond with ONLY a JSON object, no prose and no markdown fences, exactly matching:
{"task":{"title":"...","description":"...","acceptanceCriteria":["..."],"effortCoefficient":0.4},"notes":"short advice"}
Rules: effortCoefficient is the effort RELATIVE to the source task, a number in (0,2].`

func ticketPrompt(req atomizeRequest) string {
	var sb strings.Builder
	fmt.Fprintf(&sb, "TICKET: %s\n\nDESCRIPTION:\n%s\n", req.Title, req.Description)
	if len(req.AcceptanceCriteria) > 0 {
		sb.WriteString("\nACCEPTANCE CRITERIA:\n")
		for _, c := range req.AcceptanceCriteria {
			fmt.Fprintf(&sb, "- %s\n", c)
		}
	}
	if len(req.Links) > 0 {
		fmt.Fprintf(&sb, "\nLINKS: %s\n", strings.Join(req.Links, ", "))
	}
	for _, a := range req.Attachments {
		fmt.Fprintf(&sb, "ATTACHMENT: %s (%s) %s\n", a.Name, a.MimeType, a.URL)
	}
	return sb.String()
}

func (l *llmClient) atomize(ctx context.Context, req atomizeRequest, minT, maxT int) ([]subTask, string, string, error) {
	if l.cfg.apiKey == "" {
		return fallbackAtomize(req, minT, maxT), "offline-fallback", "Deterministic split (no ANTHROPIC_API_KEY configured).", nil
	}
	system := strings.NewReplacer("MIN", fmt.Sprint(minT), "MAX", fmt.Sprint(maxT)).Replace(atomizeSystem)
	user := ticketPrompt(req)
	if req.SubdivisionNote != "" {
		user += "\nCONSULTANT NOTE: " + req.SubdivisionNote
	}

	var parsed struct {
		Tasks []subTask `json:"tasks"`
		Notes string    `json:"notes"`
	}
	raw, err := l.completeWithRetry(ctx, system, user, func(text string) error {
		if err := json.Unmarshal([]byte(stripFences(text)), &parsed); err != nil {
			return err
		}
		if len(parsed.Tasks) == 0 {
			return fmt.Errorf("no tasks in response")
		}
		for _, t := range parsed.Tasks {
			if t.Title == "" || t.EffortCoefficient <= 0 {
				return fmt.Errorf("invalid task entry")
			}
		}
		return nil
	})
	if err != nil {
		l.log.Warn("llm atomize failed, using fallback", "err", err)
		return fallbackAtomize(req, minT, maxT), "offline-fallback",
			"LLM unavailable (" + err.Error() + "); deterministic split.", nil
	}
	_ = raw
	return parsed.Tasks, l.cfg.model, parsed.Notes, nil
}

func (l *llmClient) extend(ctx context.Context, req atomizeRequest) (subTask, string, string, error) {
	if l.cfg.apiKey == "" {
		return fallbackExtend(req), "offline-fallback", "Deterministic extension (no ANTHROPIC_API_KEY configured).", nil
	}
	user := ticketPrompt(req) + "\nEXTENSION NOTE (required scope): " + req.ExtensionNote

	var parsed struct {
		Task  subTask `json:"task"`
		Notes string  `json:"notes"`
	}
	_, err := l.completeWithRetry(ctx, extendSystem, user, func(text string) error {
		if err := json.Unmarshal([]byte(stripFences(text)), &parsed); err != nil {
			return err
		}
		if parsed.Task.Title == "" {
			return fmt.Errorf("no task in response")
		}
		return nil
	})
	if err != nil {
		l.log.Warn("llm extend failed, using fallback", "err", err)
		return fallbackExtend(req), "offline-fallback",
			"LLM unavailable (" + err.Error() + "); deterministic extension.", nil
	}
	return parsed.Task, l.cfg.model, parsed.Notes, nil
}

// completeWithRetry runs the chat completion and retries ONCE on parse
// failure with a corrective hint (§9.1).
func (l *llmClient) completeWithRetry(ctx context.Context, system, user string,
	validate func(string) error) (string, error) {
	text, err := l.complete(ctx, system, user)
	if err != nil {
		return "", err
	}
	if vErr := validate(text); vErr == nil {
		return text, nil
	} else {
		l.log.Warn("llm response failed validation, retrying once", "err", vErr)
	}
	text, err = l.complete(ctx, system,
		user+"\n\nIMPORTANT: your previous reply was not valid JSON matching the schema. Reply with ONLY the JSON object.")
	if err != nil {
		return "", err
	}
	if vErr := validate(text); vErr != nil {
		return "", fmt.Errorf("invalid JSON after retry: %w", vErr)
	}
	return text, nil
}

func (l *llmClient) complete(ctx context.Context, system, user string) (string, error) {
	if l.cfg.apiStyle == "anthropic" {
		return l.completeNative(ctx, system, user)
	}
	return l.completeOpenAI(ctx, system, user)
}

// completeOpenAI uses Anthropic's OpenAI-compatible endpoint.
func (l *llmClient) completeOpenAI(ctx context.Context, system, user string) (string, error) {
	body, _ := json.Marshal(map[string]any{
		"model":      l.cfg.model,
		"max_tokens": 4096,
		"messages": []map[string]string{
			{"role": "system", "content": system},
			{"role": "user", "content": user},
		},
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		l.cfg.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+l.cfg.apiKey)
	data, err := l.do(req)
	if err != nil {
		return "", err
	}
	var out struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return "", fmt.Errorf("decode chat completion: %w", err)
	}
	if len(out.Choices) == 0 {
		return "", fmt.Errorf("no choices in response")
	}
	return out.Choices[0].Message.Content, nil
}

// completeNative uses the native Anthropic Messages API (flip via
// LLM_API_STYLE=anthropic if the compatibility endpoint misbehaves, §9.1).
func (l *llmClient) completeNative(ctx context.Context, system, user string) (string, error) {
	body, _ := json.Marshal(map[string]any{
		"model":      l.cfg.model,
		"max_tokens": 4096,
		"system":     system,
		"messages":   []map[string]string{{"role": "user", "content": user}},
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		l.cfg.baseURL+"/messages", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", l.cfg.apiKey)
	req.Header.Set("anthropic-version", "2023-06-01")
	data, err := l.do(req)
	if err != nil {
		return "", err
	}
	var out struct {
		Content []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return "", fmt.Errorf("decode messages response: %w", err)
	}
	for _, c := range out.Content {
		if c.Type == "text" {
			return c.Text, nil
		}
	}
	return "", fmt.Errorf("no text content in response")
}

func (l *llmClient) do(req *http.Request) ([]byte, error) {
	resp, err := l.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(io.LimitReader(resp.Body, 10<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("llm api status %d: %.300s", resp.StatusCode, data)
	}
	return data, nil
}

// stripFences defensively removes markdown code fences and surrounding prose
// (§9.1) by slicing from the first '{' to the last '}'.
func stripFences(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.IndexByte(s, '{'); i >= 0 {
		if j := strings.LastIndexByte(s, '}'); j > i {
			return s[i : j+1]
		}
	}
	return s
}