so-many-crows/brain.go

package main

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

// Brain is whatever model the crows think with.
type Brain interface {
	Name() string
	Ask(ctx context.Context, system, prompt string) (text string, tokens int, err error)
}

// NewBrain uses Claude when ANTHROPIC_API_KEY is set, otherwise an offline mock.
func NewBrain() Brain {
	key := os.Getenv("ANTHROPIC_API_KEY")
	if key == "" {
		return MockBrain{}
	}
	model := os.Getenv("CROWS_MODEL")
	if model == "" {
		model = "claude-sonnet-5-5"
	}
	return &AnthropicBrain{key: key, model: model, client: &http.Client{}}
}

type AnthropicBrain struct {
	key, model string
	client     *http.Client
}

func (b *AnthropicBrain) Name() string { return "Claude (" + b.model + ")" }

func (b *AnthropicBrain) Ask(ctx context.Context, system, prompt string) (string, int, error) {
	body, _ := json.Marshal(map[string]any{
		"model":      b.model,
		"max_tokens": 2048,
		"system":     system,
		"messages":   []map[string]string{{"role": "user", "content": prompt}},
	})
	req, err := http.NewRequestWithContext(ctx, "POST", "https://api.anthropic.com/v1/messages", bytes.NewReader(body))
	if err != nil {
		return "", 0, err
	}
	req.Header.Set("x-api-key", b.key)
	req.Header.Set("anthropic-version", "2023-06-01")
	req.Header.Set("content-type", "application/json")

	resp, err := b.client.Do(req)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return "", 0, fmt.Errorf("model API returned %d: %s", resp.StatusCode, raw)
	}
	var out struct {
		Content []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
		Usage struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return "", 0, err
	}
	var sb strings.Builder
	for _, c := range out.Content {
		if c.Type == "text" {
			sb.WriteString(c.Text)
		}
	}
	return sb.String(), out.Usage.InputTokens + out.Usage.OutputTokens, nil
}

// MockBrain lets the whole flow run offline with no API key.
type MockBrain struct{}

func (MockBrain) Name() string { return "offline mock" }

func (MockBrain) Ask(ctx context.Context, system, prompt string) (string, int, error) {
	if err := ctx.Err(); err != nil {
		return "", 0, err
	}
	tokens := (len(system)+len(prompt))/4 + 40
	switch {
	case strings.Contains(system, markerSentry):
		return mockTriage(after(prompt, "ITEM:")), tokens, nil
	case strings.Contains(system, markerElder):
		return "[mock] Elder's reconciled answer from " + fmt.Sprint(strings.Count(prompt, "--- ")) + " reports.", tokens, nil
	default:
		return "[mock] Handled: " + after(prompt, "TASK:"), tokens, nil
	}
}

func after(s, marker string) string {
	if i := strings.LastIndex(s, marker); i >= 0 {
		return strings.TrimSpace(s[i+len(marker):])
	}
	return strings.TrimSpace(s)
}

func mockTriage(item string) string {
	l := strings.ToLower(item)
	tt := "other"
	switch {
	case hasAny(l, "find", "look up", "what is", "who ", "when "):
		tt = "retrieval"
	case hasAny(l, "write", "draft"):
		tt = "writing"
	case hasAny(l, "compare", "analyze", "should"):
		tt = "analysis"
	}
	t := Triage{Mode: "solo", TaskType: tt, Reason: "simple, one crow can handle it"}
	if parts := strings.Split(item, " and "); len(parts) > 1 {
		t = Triage{Mode: "split", TaskType: tt, Subtasks: parts, Reason: "multiple distinct parts"}
	} else if hasAny(l, "should", "decide", "best") {
		t = Triage{Mode: "murder", TaskType: tt, Reason: "judgment call, worth several opinions"}
	}
	b, _ := json.Marshal(t)
	return string(b)
}

func hasAny(s string, subs ...string) bool {
	for _, x := range subs {
		if strings.Contains(s, x) {
			return true
		}
	}
	return false
}