magpie/internal/agent/workbuddy.go

package agent

// WorkBuddy (Tencent's CodeBuddy desktop app) takes models of one's own
// from ~/.workbuddy/models.json ($WORKBUDDY_CONFIG_DIR if set): a list of
// OpenAI chat-completions models, or {"models":[…],"availableModels":[…]}
// as CodeBuddy Code's CLI writes it. Its own "local models" are entries of
// the same list (local: true).
//
//	[{"id":"deepseek/pro","name":…,"vendor":"magpie","apiKey":"magpie-workbuddy",
//	  "url":"http://127.0.0.1:3425/v1/chat/completions",
//	  "maxInputTokens":…,"maxOutputTokens":…,
//	  "supportsToolCall":true,"supportsImages":…,"supportsReasoning":true,
//	  "reasoning":{"supportedEfforts":[…],"canDisableThinking":…,"defaultEffort":…}}]
//
// WorkBuddy shows them in its picker as "custom-local:<id>", and asks url
// for model <id>, with reasoning_effort, under Bearer apiKey. It watches the
// file and picks a change up without a restart. The model is picked per
// task in WorkBuddy's window, not in a file, so what magpie sets is whether
// its models are in that picker. Its requests say "CLI/<ver> WorkBuddy/<ver>",
// so they are told apart by their key.

import (
	"bytes"
	"encoding/json"
	"os"
	"path/filepath"
	"slices"

	"github.com/yetone/magpie/internal/edit"
	"github.com/yetone/magpie/internal/gateway"
)

func workbuddy(home string) *Agent {
	dir := os.Getenv("WORKBUDDY_CONFIG_DIR")
	if dir == "" {
		dir = filepath.Join(home, ".workbuddy")
	}
	path := filepath.Join(dir, "models.json")
	return &Agent{
		ID: "workbuddy", Name: "WorkBuddy", Icon: "workbuddy-color", Aliases: []string{"work-buddy"},
		UA:  []string{"workbuddy"},
		Dir: dir, Path: path,
		Sync: func() error {
			if !workbuddyWired(path) {
				return nil
			}
			return workbuddyWrite(path, true)
		},
		Fields: []Field{{
			Key: "provider", Label: "provider",
			Get: func() string {
				if workbuddyWired(path) {
					return magpieID
				}
				return ""
			},
			Set: func(v string) error { return workbuddyWrite(path, v != "") },
			Options: func(map[string]string) []Option {
				return []Option{{Value: magpieID, Label: "magpie", Icon: "magpie", Note: "every magpie model in WorkBuddy's picker"}}
			},
		}},
	}
}

// workbuddyDoc is models.json as read: its entries, and the rest of the
// object when it is one (nil for a bare list).
type workbuddyDoc struct {
	models []json.RawMessage
	rest   map[string]json.RawMessage
}

func workbuddyRead(path string) (workbuddyDoc, error) {
	var d workbuddyDoc
	b, err := edit.Read(path)
	if err != nil || len(bytes.TrimSpace(b)) == 0 {
		return d, err
	}
	if bytes.TrimSpace(b)[0] == '[' {
		return d, json.Unmarshal(b, &d.models)
	}
	if err := json.Unmarshal(b, &d.rest); err != nil {
		return d, err
	}
	if raw, ok := d.rest["models"]; ok {
		if err := json.Unmarshal(raw, &d.models); err != nil {
			return d, err
		}
	}
	return d, nil
}

// workbuddyEntry is what magpie reads of an entry.
type workbuddyEntry struct {
	ID       string `json:"id"`
	Vendor   string `json:"vendor"`
	Disabled *bool  `json:"disabled"`
}

func workbuddyMine(raw json.RawMessage) (workbuddyEntry, bool) {
	var e workbuddyEntry
	return e, json.Unmarshal(raw, &e) == nil && e.Vendor == magpieID
}

// workbuddyWired reports whether models.json has magpie's models.
func workbuddyWired(path string) bool {
	d, _ := workbuddyRead(path)
	return slices.ContainsFunc(d.models, func(raw json.RawMessage) bool { _, ok := workbuddyMine(raw); return ok })
}

// workbuddyWrite puts magpie's models into models.json (on) or takes them
// out, leaving every other entry and key as WorkBuddy wrote it. A model
// turned off in WorkBuddy stays off.
func workbuddyWrite(path string, on bool) error {
	d, err := workbuddyRead(path)
	if err != nil {
		return err
	}
	var kept []json.RawMessage
	var ours []string
	off := map[string]bool{}
	at := -1 // where magpie's were, which they keep
	for _, raw := range d.models {
		if e, ok := workbuddyMine(raw); ok {
			if at < 0 {
				at = len(kept)
			}
			ours = append(ours, e.ID)
			if e.Disabled != nil && *e.Disabled {
				off[e.ID] = true
			}
			continue
		}
		kept = append(kept, raw)
	}
	if !on && len(ours) == 0 {
		return nil
	}
	var add []json.RawMessage
	var ids []string
	if on {
		for _, m := range magpieModels("workbuddy") {
			e := workbuddyModel(m.ID, m.Name, m.Context, maxTokens(m), m.Images, m.Efforts)
			if off[m.ID] {
				e["disabled"] = true
			}
			b, err := json.Marshal(e)
			if err != nil {
				return err
			}
			add = append(add, b)
			ids = append(ids, m.ID)
		}
	}
	if at < 0 {
		at = len(kept)
	}
	models := slices.Concat(kept[:at], add, kept[at:])
	if models == nil {
		models = []json.RawMessage{}
	}

	var out any = models
	if d.rest != nil {
		// a list of the models to show keeps showing magpie's
		if raw, ok := d.rest["availableModels"]; ok {
			var avail []string
			if json.Unmarshal(raw, &avail) == nil && len(avail) > 0 {
				avail = slices.DeleteFunc(avail, func(id string) bool { return slices.Contains(ours, id) })
				avail = append(avail, ids...)
				b, _ := json.Marshal(avail)
				d.rest["availableModels"] = b
			}
		}
		b, err := json.Marshal(models)
		if err != nil {
			return err
		}
		d.rest["models"] = b
		out = d.rest
	}
	b, err := json.MarshalIndent(out, "", "  ")
	if err != nil {
		return err
	}
	return edit.WriteAtomic(path, append(b, '\n'))
}

// workbuddyModel is one of magpie's models as a models.json entry.
func workbuddyModel(id, name string, context, output int, images bool, efforts []string) map[string]any {
	if context == 0 {
		context = 200000
	}
	e := map[string]any{
		"id": id, "name": name, "vendor": magpieID,
		"apiKey": gateway.TokenFor("workbuddy"), "url": gatewayV1() + "/chat/completions",
		"maxInputTokens":   context,
		"supportsToolCall": true, "supportsImages": images, "supportsReasoning": false,
	}
	// without it WorkBuddy caps every reply at its own default
	if out := zcodeOutput(output); out > 0 {
		e["maxOutputTokens"] = out
	}
	var levels []string
	for _, l := range efforts {
		if l != "none" && !slices.Contains(levels, l) {
			levels = append(levels, l)
		}
	}
	if len(levels) > 0 {
		e["supportsReasoning"] = true
		e["reasoning"] = map[string]any{"supportedEfforts": levels,
			"canDisableThinking": slices.Contains(efforts, "none"), "defaultEffort": zcodeDefaultLevel(levels)}
	}
	return e
}