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)
}

// modelKey lets a brain report which model actually answered a call
// (routers like openrouter/free pick a different model each time).
type modelKey struct{}

func withModelSlot(ctx context.Context, slot *string) context.Context {
	return context.WithValue(ctx, modelKey{}, slot)
}

func reportModel(ctx context.Context, model string) {
	if slot, ok := ctx.Value(modelKey{}).(*string); ok && slot != nil {
		*slot = model
	}
}

// NewBrain picks the crows' model from the environment:
//
//	CROWS_PROVIDER=openrouter  OPENROUTER_API_KEY  (default model openrouter/free)
//	CROWS_PROVIDER=anthropic   ANTHROPIC_API_KEY   (default model claude-sonnet-5-5)
//	CROWS_MODEL overrides the model for either.
//
// With no provider set, whichever key is present wins (Anthropic first).
// No key at all means the offline mock.
func NewBrain() Brain {
	provider := strings.ToLower(os.Getenv("CROWS_PROVIDER"))
	anthropicKey, openrouterKey := os.Getenv("ANTHROPIC_API_KEY"), os.Getenv("OPENROUTER_API_KEY")
	if provider == "" {
		switch {
		case anthropicKey != "":
			provider = "anthropic"
		case openrouterKey != "":
			provider = "openrouter"
		}
	}
	model := os.Getenv("CROWS_MODEL")
	switch provider {
	case "anthropic":
		if anthropicKey == "" {
			fmt.Fprintln(os.Stderr, "CROWS_PROVIDER=anthropic but ANTHROPIC_API_KEY is empty; using the offline mock")
			return MockBrain{}
		}
		if model == "" {
			model = "claude-sonnet-5-5"
		}
		return &AnthropicBrain{key: anthropicKey, model: model, client: &http.Client{}}
	case "openrouter":
		if openrouterKey == "" {
			fmt.Fprintln(os.Stderr, "CROWS_PROVIDER=openrouter but OPENROUTER_API_KEY is empty; using the offline mock")
			return MockBrain{}
		}
		if model == "" {
			model = "openrouter/free"
		}
		return &ChatBrain{
			provider: "OpenRouter", url: "https://openrouter.ai/api/v1/chat/completions",
			key: openrouterKey, model: model, client: &http.Client{},
			headers: map[string]string{"X-OpenRouter-Title": "So Many Crows"},
		}
	case "":
		return MockBrain{}
	default:
		fmt.Fprintf(os.Stderr, "unknown CROWS_PROVIDER %q; using the offline mock\n", provider)
		return MockBrain{}
	}
}

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)
		}
	}
	reportModel(ctx, b.model)
	return sb.String(), out.Usage.InputTokens + out.Usage.OutputTokens, nil
}

// ChatBrain speaks the OpenAI-style chat completions format (OpenRouter, and
// gateways like OpenCode Zen for open models).
type ChatBrain struct {
	provider, url, key, model string
	headers                   map[string]string
	client                    *http.Client
}

func (b *ChatBrain) Name() string { return b.provider + " (" + b.model + ")" }

func (b *ChatBrain) Ask(ctx context.Context, system, prompt string) (string, int, error) {
	body, _ := json.Marshal(map[string]any{
		"model":      b.model,
		"max_tokens": 2048,
		"messages": []map[string]string{
			{"role": "system", "content": system},
			{"role": "user", "content": prompt},
		},
	})
	req, err := http.NewRequestWithContext(ctx, "POST", b.url, bytes.NewReader(body))
	if err != nil {
		return "", 0, err
	}
	req.Header.Set("Authorization", "Bearer "+b.key)
	req.Header.Set("Content-Type", "application/json")
	for k, v := range b.headers {
		req.Header.Set(k, v)
	}
	resp, err := b.client.Do(req)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if resp.StatusCode != http.StatusOK {
		return "", 0, fmt.Errorf("%s returned %d: %s", b.provider, resp.StatusCode, clip(raw, 300))
	}
	var out struct {
		Model   string `json:"model"`
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage *struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
		} `json:"usage"`
		Error *struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return "", 0, fmt.Errorf("%s sent an unreadable reply: %w", b.provider, err)
	}
	if out.Error != nil {
		return "", 0, fmt.Errorf("%s error: %s", b.provider, out.Error.Message)
	}
	if len(out.Choices) == 0 || strings.TrimSpace(out.Choices[0].Message.Content) == "" {
		return "", 0, fmt.Errorf("%s sent an empty reply", b.provider)
	}
	text := out.Choices[0].Message.Content
	// Budgets must still bite if a provider leaves out usage: estimate ~4 characters per token.
	tokens := (len(system) + len(prompt) + len(text)) / 4
	if out.Usage != nil && out.Usage.PromptTokens+out.Usage.CompletionTokens > 0 {
		tokens = out.Usage.PromptTokens + out.Usage.CompletionTokens
	}
	model := out.Model
	if model == "" {
		model = b.model
	}
	reportModel(ctx, model)
	return text, tokens, nil
}

func clip(b []byte, n int) string {
	s := strings.TrimSpace(string(b))
	if len(s) > n {
		return s[:n] + "…"
	}
	return s
}

// 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
	}
	reportModel(ctx, "mock")
	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
}