Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d1e91868ee |
No files matched your search
+6
-11
@@ -28,17 +28,12 @@ APIs change between minors without a deprecation period.
|
||||
|
||||
### Added
|
||||
|
||||
- **Pictures, 3D models, vision and workflows as MCP tools** at `/mcp`,
|
||||
beside the memory notes, each offered only where the gateway has a model
|
||||
for it: `generate_image` (with `concept`, `side` and `texture` styles),
|
||||
`generate_mesh` (from a picture or a prompt), `review_images` (a vision
|
||||
model, to a JSON schema), `run_workflow`, `job_status` and
|
||||
`gateway_status`. Each makes its request under the caller's key, so
|
||||
allowlists, limits, licences and budgets apply and it is charged like
|
||||
any request; what it makes is kept as the caller's file and answered by
|
||||
id and download URL. A client that takes a stream gets
|
||||
`notifications/progress` while a mesh is made and waits up to ten
|
||||
minutes; one that does not, 50 seconds, then the job id.
|
||||
- **Setup asks where alerts go**: `llmgw server init --webhook URL`
|
||||
(repeatable) and a **Notification webhooks** field on both setup screens.
|
||||
Slack's, Mattermost's and Rocket.Chat's addresses are sent Slack's text,
|
||||
any other the event as JSON (`slack:` or `json:` in front says which).
|
||||
They are added to the cluster's `notify.webhooks` beside those already
|
||||
there, and each is sent a test once the control plane runs on them.
|
||||
|
||||
- **Extended thinking on `/v1/messages`**: `thinking: {type: "enabled"}`
|
||||
tells a model that thinks by its chat template to (`enable_thinking`),
|
||||
|
||||
@@ -82,7 +82,7 @@
|
||||
"/mcp": {
|
||||
"post": {
|
||||
"operationId": "postMcp",
|
||||
"summary": "The gateway's MCP server: the memory notes, and pictures, 3D models, vision and workflows, as tools",
|
||||
"summary": "The gateway's MCP server: the memory notes as tools",
|
||||
"tags": [
|
||||
"mcp"
|
||||
],
|
||||
|
||||
@@ -534,42 +534,6 @@ curl https://llm.example.com/mcp \
|
||||
-d '{"jsonrpc": "2.0", "id": 1, "method": "tools/list"}'
|
||||
```
|
||||
|
||||
**Pictures, 3D models, vision and workflows.** Where the gateway has the
|
||||
models for them, the same server offers more tools — only those it has a
|
||||
model for:
|
||||
|
||||
| Tool | Does |
|
||||
|---|---|
|
||||
| `generate_image` | A picture from a prompt. `style` adds what a picture for a mesh wants: `concept` (one object, three-quarter view, grey background), `side` (an exact side view), `texture` (a seamless tile) or `none`. Answers the PNG's `file_id`, `url` and `seed`, with a 512 px preview to look at. |
|
||||
| `generate_mesh` | A GLB from a picture — `image` is a `file-…` id (a `generate_image` result, say), a data URL or base64 — or from a `prompt`, for which a concept picture is made first. `seed`, `resolution`, `texture_resolution`, and any other input the model's workflow declares in `inputs`. Answers the GLB's `file_id`, `url` and licence. |
|
||||
| `review_images` | Asks a vision model about one to four pictures (each shrunk to 1280 px on its long side); `json_schema` makes it answer that JSON. |
|
||||
| `run_workflow` | Runs a ComfyUI workflow from the inputs its catalog entry declares, and answers its result file. |
|
||||
| `job_status` | Where a mesh or a workflow run stands, and its file when done; `wait_seconds` waits for it. |
|
||||
| `gateway_status` | The picture, 3D, vision and workflow models, whether each is loaded on a GPU now, a workflow's inputs, and your last jobs. |
|
||||
|
||||
Each tool makes its request of the gateway under your key, so your
|
||||
allowlist, limits, left-out licences and budget apply, and it is charged
|
||||
and listed like any other request. What it makes is kept as your file for
|
||||
as long as job results are: fetch it from `url` with the same key (`curl -H
|
||||
"Authorization: Bearer $LLMGW_API_KEY" -o rock.glb <url>`). A tool server
|
||||
over HTTP cannot write to your disk, so an agent that needs the file
|
||||
locally — to import it in Blender — fetches it, or goes through a small
|
||||
local bridge that does.
|
||||
|
||||
A mesh takes from half a minute to a few minutes. A client that takes the
|
||||
answer as a stream (`Accept: text/event-stream`, as Claude Code does) is
|
||||
answered that way: `notifications/progress` for the call's
|
||||
`progressToken` while it runs — queued, waiting for the GPU, the node and
|
||||
its percent — then the result, and the call waits up to ten minutes. One
|
||||
answered in one piece waits 50 seconds; a job not done by then is answered
|
||||
with its `job_id`, to ask after with `job_status`. `wait_seconds` (up to
|
||||
900) sets either.
|
||||
|
||||
```
|
||||
claude mcp add --transport http llmgw https://llm.example.com/mcp \
|
||||
--header "Authorization: Bearer $LLMGW_API_KEY"
|
||||
```
|
||||
|
||||
### Licences you leave out
|
||||
|
||||
A model's licence is in `/v1/models` (`license`: its `id`, `name`, whether
|
||||
|
||||
@@ -1827,6 +1827,23 @@ same event about the same thing is sent at most once an hour, so a budget that
|
||||
refuses a thousand requests is one message. `llmgw ctl notify test` sends a
|
||||
test event to each webhook and says how each answered.
|
||||
|
||||
Setup asks for them: **Notification webhooks** on its first page, or
|
||||
`--webhook` to `llmgw server init`, as often as there are webhooks:
|
||||
|
||||
```
|
||||
sudo llmgw server init --admin ops@example.com \
|
||||
--webhook https://hooks.slack.com/services/… \
|
||||
--webhook json:https://ops.example.com/llmgw
|
||||
```
|
||||
|
||||
An address of Slack's, or one with `/hooks/` in it (Mattermost, Rocket.Chat)
|
||||
or ending in `/slack`, is sent Slack's `{"text": …}`; any other the event as
|
||||
JSON; `slack:` or `json:` in front says which. They are added to the
|
||||
cluster's settings beside any already there — run again, nothing is added
|
||||
twice — and once the control plane runs on them, each is sent a test, whose
|
||||
answer setup prints. Events to choose and a signing secret are set in the
|
||||
settings afterwards (`llmgw ctl config edit`).
|
||||
|
||||
### Tracing
|
||||
|
||||
Metrics say that a model took four minutes to start; a trace says where they
|
||||
|
||||
@@ -73,6 +73,9 @@ type serverInit struct {
|
||||
// patroni are user@host of the machines to keep the store on under
|
||||
// Patroni, replicated; this machine is the third etcd vote for two.
|
||||
patroni []string
|
||||
// webhooks are where the cluster's alerts go: put into its settings
|
||||
// once it runs, and each sent a test.
|
||||
webhooks []string
|
||||
// patroniNodes are the same machines already connected to, which the
|
||||
// setup screen does with the password it was given; when set, patroni
|
||||
// only names them.
|
||||
@@ -182,6 +185,9 @@ func newServerInitCmd() *cobra.Command {
|
||||
"leave the firewall alone, instead of keeping llmgw's internal ports to the cluster's networks")
|
||||
f.BoolVar(&in.postgres, "postgres", false,
|
||||
"keep the store in PostgreSQL on this machine, installing it and moving an existing SQLite store into it")
|
||||
f.StringArrayVar(&in.webhooks, "webhook", nil,
|
||||
"where alerts go — a model parked, a node lost, a budget spent: an http(s) address, Slack's, Mattermost's "+
|
||||
"and Rocket.Chat's sent Slack's text, any other the event as JSON (slack: or json: in front says which); repeatable")
|
||||
f.StringArrayVar(&in.patroni, "patroni", nil,
|
||||
"user@host of a machine to keep the store on, replicated under Patroni; give two or more. "+
|
||||
"Installs PostgreSQL, Patroni and etcd there over ssh, and moves the existing store in")
|
||||
@@ -249,6 +255,9 @@ func (in *serverInit) run(ctx context.Context) error {
|
||||
fmt.Fprintf(in.out, " !! firewall: %v; llmgw's ports answer every network this machine is on\n", err)
|
||||
}
|
||||
}
|
||||
if err := in.ensureWebhooks(ctx); err != nil {
|
||||
fmt.Fprintf(in.out, " !! webhooks: %v\n", err)
|
||||
}
|
||||
if in.judgeWanted() && in.configJudgesWith(guardModelName) {
|
||||
// Not fatal: without its judge the baseline policy still checks for
|
||||
// secrets and injection, and a judged request is served flagged.
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.yaml.in/yaml/v3"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
|
||||
)
|
||||
|
||||
// Where the cluster says what went wrong: a model parked, a node lost, a
|
||||
// budget spent. notify.webhooks existed, and setup never asked for one, so a
|
||||
// site learnt of a lost node from its users.
|
||||
|
||||
// parseWebhook reads a webhook as setup takes it: an address, with "slack:"
|
||||
// or "json:" in front to say what it is sent, else what its address says.
|
||||
func parseWebhook(s string) (notify.Webhook, error) {
|
||||
s = strings.TrimSpace(s)
|
||||
format := ""
|
||||
for _, f := range []string{"slack", "json"} {
|
||||
if rest, ok := strings.CutPrefix(s, f+":"); ok && !strings.HasPrefix(rest, "//") {
|
||||
format, s = f, rest
|
||||
}
|
||||
}
|
||||
u, err := url.Parse(s)
|
||||
if err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Host == "" {
|
||||
return notify.Webhook{}, fmt.Errorf("webhook %q: an http(s) address, with slack: or json: in front to say its format", s)
|
||||
}
|
||||
if format == "" {
|
||||
format = webhookFormat(u)
|
||||
}
|
||||
return notify.Webhook{URL: s, Format: format}, nil
|
||||
}
|
||||
|
||||
// webhookFormat is what an address takes: Slack's own, and the incoming
|
||||
// webhooks that read Slack's {"text": …} — Mattermost's and Rocket.Chat's
|
||||
// /hooks/, Discord's …/slack — are sent that; anything else the event as
|
||||
// JSON.
|
||||
func webhookFormat(u *url.URL) string {
|
||||
switch {
|
||||
case u.Host == "hooks.slack.com",
|
||||
strings.Contains(u.Path, "/hooks/"),
|
||||
strings.HasSuffix(u.Path, "/slack"):
|
||||
return "slack"
|
||||
}
|
||||
return "json"
|
||||
}
|
||||
|
||||
// parseWebhooks reads a list of them, by commas or as given.
|
||||
func parseWebhooks(list []string) ([]notify.Webhook, error) {
|
||||
var out []notify.Webhook
|
||||
for _, item := range list {
|
||||
for _, s := range strings.Split(item, ",") {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
continue
|
||||
}
|
||||
w, err := parseWebhook(s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, w)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// webhookEntry is one of notify.webhooks as the settings document has it.
|
||||
type webhookEntry struct {
|
||||
URL string `yaml:"url"`
|
||||
Format string `yaml:"format,omitempty"`
|
||||
Events []string `yaml:"events,omitempty,flow"`
|
||||
SecretFile string `yaml:"secret_file,omitempty"`
|
||||
}
|
||||
|
||||
// withWebhooks adds these webhooks to a settings document's notify.webhooks,
|
||||
// and says whether that changed it: one already there, by address, is left
|
||||
// as it is, with whatever events and secret it was given.
|
||||
func withWebhooks(doc *settings.Doc, add []notify.Webhook) (bool, error) {
|
||||
path := []string{"notify", "webhooks"}
|
||||
var list []webhookEntry
|
||||
if n := doc.Node(path); n != nil {
|
||||
if err := n.Decode(&list); err != nil {
|
||||
return false, fmt.Errorf("notify.webhooks in the cluster's settings: %w", err)
|
||||
}
|
||||
}
|
||||
changed := false
|
||||
for _, w := range add {
|
||||
if !slices.ContainsFunc(list, func(e webhookEntry) bool { return e.URL == w.URL }) {
|
||||
list = append(list, webhookEntry{URL: w.URL, Format: w.Format})
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return false, nil
|
||||
}
|
||||
b, err := yaml.Marshal(list)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, doc.SetYAML(path, string(b))
|
||||
}
|
||||
|
||||
// ensureWebhooks puts the webhooks given into the cluster's settings and
|
||||
// sends each a test, once the control plane runs on them. Said, not fatal: a
|
||||
// cluster without its alerts still serves.
|
||||
func (in *serverInit) ensureWebhooks(ctx context.Context) error {
|
||||
if len(in.webhooks) == 0 {
|
||||
return nil
|
||||
}
|
||||
hooks, err := parseWebhooks(in.webhooks)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c := in.adminAPI()
|
||||
var cur adminapi.SettingsView
|
||||
if err := c.get(ctx, "/api/v1/settings", &cur); err != nil {
|
||||
return err
|
||||
}
|
||||
doc, err := settings.FromBytes("settings", []byte(cur.Body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
changed, err := withWebhooks(doc, hooks)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if changed {
|
||||
b, err := doc.Bytes()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req, _ := json.Marshal(adminapi.SettingsRequest{Base: cur.Version, Body: string(b), Note: "notification webhooks, from setup"})
|
||||
if err := c.put(ctx, "/api/v1/settings", req, &cur); err != nil {
|
||||
return err
|
||||
}
|
||||
in.done("notification webhooks in the cluster's settings (version %d)", cur.Version)
|
||||
} else {
|
||||
in.kept("notification webhooks (already in the cluster's settings)")
|
||||
}
|
||||
|
||||
// The control planes start again on the change; the test is sent by
|
||||
// one that knows every webhook asked for.
|
||||
var res []notify.Result
|
||||
err = waitFor(ctx, 2*time.Minute, "the control plane to run on its webhooks", func() (bool, error) {
|
||||
res = nil
|
||||
if c.post(ctx, "/api/v1/notify/test", nil, &res) != nil {
|
||||
return false, nil
|
||||
}
|
||||
for _, w := range hooks {
|
||||
// Answered by host: the path of a webhook is its secret.
|
||||
if !slices.ContainsFunc(res, func(r notify.Result) bool { return r.URL == notify.Redact(w.URL) }) {
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var failed []string
|
||||
for _, r := range res {
|
||||
if !slices.ContainsFunc(hooks, func(w notify.Webhook) bool { return notify.Redact(w.URL) == r.URL }) {
|
||||
continue
|
||||
}
|
||||
if r.Err != "" {
|
||||
failed = append(failed, r.URL+": "+r.Err)
|
||||
continue
|
||||
}
|
||||
in.done("a test notification reached %s", r.URL)
|
||||
}
|
||||
if len(failed) > 0 {
|
||||
return errors.New("a test notification did not arrive: " + strings.Join(failed, "; ") +
|
||||
" (kept in the settings: 'llmgw ctl notify test' tries again)")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
|
||||
)
|
||||
|
||||
func TestAWebhooksFormatIsWhatItsAddressTakes(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"https://hooks.slack.com/services/T1/B2/abc": "slack",
|
||||
"https://chat.example.com/hooks/xyz": "slack", // Mattermost, Rocket.Chat
|
||||
"https://discord.com/api/webhooks/1/abc/slack": "slack",
|
||||
"https://alerts.example.com/llmgw": "json",
|
||||
"json:https://hooks.slack.com/services/T1/B2/abc": "json",
|
||||
"slack:https://alerts.example.com/llmgw": "slack",
|
||||
" https://events.pagerduty.com/integration/x/enqueue ": "json",
|
||||
} {
|
||||
w, err := parseWebhook(in)
|
||||
if err != nil || w.Format != want || strings.HasPrefix(w.URL, "slack:") || strings.HasPrefix(w.URL, "json:") {
|
||||
t.Errorf("%q: %+v %v, want %s", in, w, err, want)
|
||||
}
|
||||
}
|
||||
for _, bad := range []string{"hooks.slack.com/x", "ftp://example.com/x", "slack:", "https://"} {
|
||||
if _, err := parseWebhook(bad); err == nil {
|
||||
t.Errorf("%q was taken", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebhooksAreAddedToTheSettingsBesideTheOnesThere(t *testing.T) {
|
||||
doc, err := settings.FromBytes("settings", []byte("router:\n retries: 2\nnotify:\n webhooks:\n"+
|
||||
" - url: https://alerts.example.com/a\n events: [node_lost]\n secret_file: /etc/llmgw/secrets/hook\n"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hooks, _ := parseWebhooks([]string{"https://alerts.example.com/a, https://hooks.slack.com/services/T/B/c"})
|
||||
changed, err := withWebhooks(doc, hooks)
|
||||
if err != nil || !changed {
|
||||
t.Fatalf("changed %v: %v", changed, err)
|
||||
}
|
||||
b, _ := doc.Bytes()
|
||||
got := string(b)
|
||||
for _, want := range []string{"retries: 2", "url: https://alerts.example.com/a", "events: [node_lost]",
|
||||
"secret_file: /etc/llmgw/secrets/hook", "url: https://hooks.slack.com/services/T/B/c", "format: slack"} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Errorf("%q not in:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
if strings.Count(got, "alerts.example.com/a") != 1 {
|
||||
t.Errorf("a webhook already there was added again:\n%s", got)
|
||||
}
|
||||
if changed, _ := withWebhooks(doc, hooks); changed {
|
||||
t.Error("adding the same webhooks again changed the settings")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetupPutsItsWebhooksInTheSettingsAndTestsThemOnceInForce(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
body, tests := "router:\n retries: 2\n", 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch {
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/api/v1/settings":
|
||||
_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 1, Body: body})
|
||||
case r.Method == http.MethodPut && r.URL.Path == "/api/v1/settings":
|
||||
var req adminapi.SettingsRequest
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
body = req.Body
|
||||
_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 2, Body: body})
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/notify/test":
|
||||
tests++
|
||||
// The first answer comes from the control plane before it
|
||||
// started again on the change: it knows no webhooks.
|
||||
res := []notify.Result{}
|
||||
if tests > 1 {
|
||||
res = []notify.Result{{URL: notify.Redact("https://hooks.slack.com/services/T/B/c")},
|
||||
{URL: notify.Redact("https://ops.example.com/in"), Err: "502 Bad Gateway"}}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(res)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "server.yaml"),
|
||||
[]byte("admin_listen: "+strings.TrimPrefix(srv.URL, "http://")+"\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out bytes.Buffer
|
||||
in := &serverInit{etcDir: dir, out: &out,
|
||||
webhooks: []string{"https://hooks.slack.com/services/T/B/c", "https://ops.example.com/in"}}
|
||||
err := in.ensureWebhooks(t.Context())
|
||||
if err == nil || !strings.Contains(err.Error(), "502") || !strings.Contains(err.Error(), "ctl notify test") {
|
||||
t.Errorf("a webhook that did not answer: %v", err)
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if !strings.Contains(body, "format: slack") || !strings.Contains(body, "url: https://ops.example.com/in") ||
|
||||
!strings.Contains(body, "retries: 2") || tests < 2 {
|
||||
t.Errorf("settings %q after %d tests", body, tests)
|
||||
}
|
||||
if strings.Contains(out.String(), "/services/T/B/c") {
|
||||
t.Errorf("a webhook's secret path was printed:\n%s", out.String())
|
||||
}
|
||||
if !strings.Contains(out.String(), "reached https://hooks.slack.com/…") {
|
||||
t.Errorf("the one that answered is not said:\n%s", out.String())
|
||||
}
|
||||
}
|
||||
@@ -50,6 +50,7 @@ func (w *lineWriter) Write(p []byte) (int, error) {
|
||||
const (
|
||||
fAdmin = iota
|
||||
fHosts
|
||||
fWebhooks
|
||||
fPostgres
|
||||
fMirror
|
||||
fNode
|
||||
@@ -70,6 +71,7 @@ type initModel struct {
|
||||
focus int
|
||||
admin textField
|
||||
hosts textField
|
||||
webhooks textField
|
||||
bundle textField
|
||||
bundleSHA textField
|
||||
toggles [fCount]bool
|
||||
@@ -103,6 +105,11 @@ func newInitModel(in *serverInit, found []string, askForm bool) initModel {
|
||||
hosts.SetValue(strings.Join(in.hosts, ", "))
|
||||
hosts.CharLimit = 400
|
||||
hosts.Width = 40
|
||||
webhooks := newTextField()
|
||||
webhooks.Placeholder = "https://hooks.slack.com/services/…"
|
||||
webhooks.SetValue(strings.Join(in.webhooks, ", "))
|
||||
webhooks.CharLimit = 800
|
||||
webhooks.Width = 40
|
||||
bundle := newTextField()
|
||||
bundle.Placeholder = "/media/usb/llmgw-bundle.tar"
|
||||
bundle.SetValue(in.bundle)
|
||||
@@ -114,7 +121,7 @@ func newInitModel(in *serverInit, found []string, askForm bool) initModel {
|
||||
bundleSHA.CharLimit = 64
|
||||
bundleSHA.Width = 64
|
||||
sp := newSpinner()
|
||||
m := initModel{in: in, found: found, form: askForm, admin: admin, hosts: hosts,
|
||||
m := initModel{in: in, found: found, form: askForm, admin: admin, hosts: hosts, webhooks: webhooks,
|
||||
bundle: bundle, bundleSHA: bundleSHA, spin: sp, width: 100}
|
||||
m.toggles[fPostgres] = in.postgres
|
||||
m.toggles[fMirror] = !in.noMirror
|
||||
@@ -235,6 +242,17 @@ func (m initModel) formKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.moveFocus(0)
|
||||
return m, nil
|
||||
}
|
||||
hooks := strings.TrimSpace(m.webhooks.Value())
|
||||
if _, err := parseWebhooks([]string{hooks}); err != nil {
|
||||
m.formError = err.Error()
|
||||
m.focus = fWebhooks
|
||||
m.moveFocus(0)
|
||||
return m, nil
|
||||
}
|
||||
m.in.webhooks = nil
|
||||
if hooks != "" {
|
||||
m.in.webhooks = []string{hooks}
|
||||
}
|
||||
m.in.bundle, m.in.bundleSHA = bundle, sum
|
||||
m.in.admin = admin
|
||||
m.in.hosts = nil
|
||||
@@ -260,6 +278,10 @@ func (m initModel) formKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
var cmd tea.Cmd
|
||||
m.hosts, cmd = m.hosts.Update(k)
|
||||
return m, cmd
|
||||
case fWebhooks:
|
||||
var cmd tea.Cmd
|
||||
m.webhooks, cmd = m.webhooks.Update(k)
|
||||
return m, cmd
|
||||
case fBundle:
|
||||
var cmd tea.Cmd
|
||||
m.bundle, cmd = m.bundle.Update(k)
|
||||
@@ -284,6 +306,7 @@ func (m *initModel) moveFocus(d int) {
|
||||
}
|
||||
m.admin.Blur()
|
||||
m.hosts.Blur()
|
||||
m.webhooks.Blur()
|
||||
m.bundle.Blur()
|
||||
m.bundleSHA.Blur()
|
||||
switch m.focus {
|
||||
@@ -291,6 +314,8 @@ func (m *initModel) moveFocus(d int) {
|
||||
m.admin.Focus()
|
||||
case fHosts:
|
||||
m.hosts.Focus()
|
||||
case fWebhooks:
|
||||
m.webhooks.Focus()
|
||||
case fBundle:
|
||||
m.bundle.Focus()
|
||||
case fBundleSHA:
|
||||
@@ -409,6 +434,8 @@ func (m initModel) formView() string {
|
||||
found = strings.Join(m.found, ", ")
|
||||
}
|
||||
b.WriteString(" " + sDim.Render("Already in the certificate: "+found+". Add a DNS name or NAT address here.") + "\n\n")
|
||||
fmt.Fprintf(&b, "%s\n %s\n", label(fWebhooks, "Notification webhooks"), m.webhooks.View())
|
||||
b.WriteString(" " + sDim.Render("Where a parked model, a lost node or a spent budget is said: Slack, Mattermost, or any URL taking JSON. Comma-separated; a test is sent.") + "\n\n")
|
||||
b.WriteString(check(fPostgres, "PostgreSQL", "rather than SQLite; installed here, for several control planes or a larger store"))
|
||||
if m.in.mirrorURL != "" {
|
||||
b.WriteString(" Model mirror " + sDim.Render("at "+m.in.mirrorURL+", from --mirror") + "\n")
|
||||
|
||||
@@ -20,6 +20,8 @@ func TestTheInitFormBecomesInitsSettings(t *testing.T) {
|
||||
m = typeIn(m, "ops@example.com")
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
|
||||
m = typeIn(m, "llm.example.com, 203.0.113.10")
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
|
||||
m = typeIn(m, "https://hooks.slack.com/services/T1/B2/x, json:https://alerts.example.com/in")
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(" ")}) // PostgreSQL on
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab},
|
||||
tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(" ")}) // starter model off
|
||||
@@ -33,11 +35,30 @@ func TestTheInitFormBecomesInitsSettings(t *testing.T) {
|
||||
if !in.postgres || in.noMirror || in.noNode || !in.noModel || !in.noGuardModel {
|
||||
t.Errorf("toggles: postgres %v mirror %v node %v model %v", in.postgres, !in.noMirror, !in.noNode, !in.noModel)
|
||||
}
|
||||
if hooks, err := parseWebhooks(in.webhooks); err != nil || len(hooks) != 2 || hooks[0].Format != "slack" || hooks[1].Format != "json" {
|
||||
t.Errorf("webhooks %q: %+v %v", in.webhooks, hooks, err)
|
||||
}
|
||||
if m.(initModel).form {
|
||||
t.Error("the form stayed up after Set up")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWebhookThatIsNoAddressKeepsTheFormUp(t *testing.T) {
|
||||
in := &serverInit{}
|
||||
var m tea.Model = newInitModel(in, nil, true)
|
||||
for range fWebhooks {
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
|
||||
}
|
||||
m = typeIn(m, "hooks.slack.com/services/x")
|
||||
for range fStart - fWebhooks {
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
|
||||
}
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyEnter})
|
||||
if !m.(initModel).form || m.(initModel).focus != fWebhooks || !strings.Contains(m.View(), "http(s) address") {
|
||||
t.Errorf("an address without its scheme was taken:\n%s", m.View())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnOfflineBundleIsGivenInTheFormWithItsDigest(t *testing.T) {
|
||||
in := &serverInit{}
|
||||
var m tea.Model = newInitModel(in, nil, true)
|
||||
@@ -117,9 +138,10 @@ func TestWithAMirrorElsewhereTheFormDoesNotOfferOne(t *testing.T) {
|
||||
if v := m.View(); !strings.Contains(v, "at http://10.0.0.150:8082, from --mirror") {
|
||||
t.Errorf("the mirror in use is not shown:\n%s", v)
|
||||
}
|
||||
// Tab from the addresses goes past the mirror line to PostgreSQL, then
|
||||
// the node.
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab})
|
||||
// Tab from the addresses goes to the webhooks, past the mirror line to
|
||||
// PostgreSQL, then the node.
|
||||
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab},
|
||||
tea.KeyMsg{Type: tea.KeyTab})
|
||||
if f := m.(initModel).focus; f != fNode {
|
||||
t.Errorf("focus %d, want the node toggle %d", f, fNode)
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
|
||||
joinpkg "gitea.mynodes.xyz/ilker/llmgw/internal/join"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/preflight"
|
||||
@@ -150,8 +151,8 @@ type setupModel struct {
|
||||
ips []string
|
||||
choice int // 0 one machine, 1 high availability
|
||||
|
||||
admin, names textField
|
||||
afocus int // admin, names, next
|
||||
admin, names, webhooks textField
|
||||
afocus int // admin, names, webhooks, next
|
||||
|
||||
db [2]dbMachine
|
||||
mfocus int // over both machines' fields, then the check button
|
||||
@@ -184,8 +185,10 @@ func newSetupModel(host string, ips []string) setupModel {
|
||||
admin.Placeholder, admin.CharLimit, admin.Width = "you@example.com", 200, 40
|
||||
names := newTextField()
|
||||
names.Placeholder, names.CharLimit, names.Width = "llm.example.com, 203.0.113.10", 400, 40
|
||||
webhooks := newTextField()
|
||||
webhooks.Placeholder, webhooks.CharLimit, webhooks.Width = "https://hooks.slack.com/services/…", 800, 40
|
||||
return setupModel{
|
||||
ref: &progRef{}, host: host, ips: ips, admin: admin, names: names,
|
||||
ref: &progRef{}, host: host, ips: ips, admin: admin, names: names, webhooks: webhooks,
|
||||
db: [2]dbMachine{newDBMachine(1), newDBMachine(2)}, spin: newSpinner(), width: 100, height: 30,
|
||||
log: &[]string{}, logMu: &sync.Mutex{},
|
||||
}
|
||||
@@ -308,14 +311,17 @@ func (m setupModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
|
||||
func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
move := func(d int) {
|
||||
m.afocus = (m.afocus + d + 3) % 3
|
||||
m.afocus = (m.afocus + d + 4) % 4
|
||||
m.admin.Blur()
|
||||
m.names.Blur()
|
||||
m.webhooks.Blur()
|
||||
switch m.afocus {
|
||||
case 0:
|
||||
m.admin.Focus()
|
||||
case 1:
|
||||
m.names.Focus()
|
||||
case 2:
|
||||
m.webhooks.Focus()
|
||||
}
|
||||
}
|
||||
switch k.Type {
|
||||
@@ -326,7 +332,7 @@ func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
move(-1)
|
||||
return m, nil
|
||||
case tea.KeyEnter:
|
||||
if m.afocus < 2 {
|
||||
if m.afocus < 3 {
|
||||
move(1)
|
||||
return m, nil
|
||||
}
|
||||
@@ -336,6 +342,12 @@ func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
move(0)
|
||||
return m, nil
|
||||
}
|
||||
if _, err := parseWebhooks(m.webhookList()); err != nil {
|
||||
m.formErr = err.Error()
|
||||
m.afocus = 2
|
||||
move(0)
|
||||
return m, nil
|
||||
}
|
||||
m.formErr = ""
|
||||
m.page, m.mfocus = pMachines, 0
|
||||
m.db[0].addr.Focus()
|
||||
@@ -347,10 +359,23 @@ func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.admin, cmd = m.admin.Update(k)
|
||||
case 1:
|
||||
m.names, cmd = m.names.Update(k)
|
||||
case 2:
|
||||
m.webhooks, cmd = m.webhooks.Update(k)
|
||||
}
|
||||
return m, cmd
|
||||
}
|
||||
|
||||
// webhookList is the webhooks asked for, one an item.
|
||||
func (m setupModel) webhookList() []string {
|
||||
var out []string
|
||||
for _, s := range strings.Split(m.webhooks.Value(), ",") {
|
||||
if s = strings.TrimSpace(s); s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (m *setupModel) focusMachine() {
|
||||
for i := range m.db {
|
||||
for f := 0; f < mfCount; f++ {
|
||||
@@ -589,6 +614,9 @@ func (m setupModel) command() string {
|
||||
for _, t := range m.targets() {
|
||||
parts = append(parts, "--patroni "+t)
|
||||
}
|
||||
for _, w := range m.webhookList() {
|
||||
parts = append(parts, "--webhook "+w)
|
||||
}
|
||||
parts = append(parts, "--no-node", "--no-model")
|
||||
return strings.Join(parts, " \\\n ")
|
||||
}
|
||||
@@ -600,6 +628,7 @@ func (m *setupModel) startRun() tea.Cmd {
|
||||
in.admin = strings.TrimSpace(m.admin.Value())
|
||||
in.hosts = m.namesList()
|
||||
in.patroni = m.targets()
|
||||
in.webhooks = m.webhookList()
|
||||
// The control plane of a replicated setup is not where models run;
|
||||
// nodes are added once it is up.
|
||||
in.noNode, in.noModel = true, true
|
||||
@@ -759,7 +788,9 @@ func (m setupModel) adminView() string {
|
||||
sDim.Render("A password is generated and shown once, at the end."))
|
||||
fmt.Fprintf(&b, "%s\n %s\n %s\n\n", label(1, "Other names nodes will reach this machine by"), m.names.View(),
|
||||
sDim.Render("Known: "+strings.Join(append([]string{m.host}, m.ips...), ", ")+". Add a DNS name or a floating address."))
|
||||
b.WriteString(button(m.afocus == 2, "Next") + "\n")
|
||||
fmt.Fprintf(&b, "%s\n %s\n %s\n\n", label(2, "Notification webhooks"), m.webhooks.View(),
|
||||
sDim.Render("Where a parked model, a lost node or a spent budget is said: Slack, Mattermost, or any URL taking JSON. Optional; a test is sent."))
|
||||
b.WriteString(button(m.afocus == 3, "Next") + "\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
@@ -880,6 +911,13 @@ func (m setupModel) planView() string {
|
||||
fmt.Fprintf(&b, " %s %s\n %s\n", lipgloss.NewStyle().Bold(true).Render(m.host), sDim.Render(m.thisAddr()),
|
||||
"control plane and gateway, model mirror, etcd (the third vote)")
|
||||
fmt.Fprintf(&b, " administrator %s\n", strings.TrimSpace(m.admin.Value()))
|
||||
if hooks, _ := parseWebhooks(m.webhookList()); len(hooks) > 0 {
|
||||
var where []string
|
||||
for _, h := range hooks {
|
||||
where = append(where, notify.Redact(h.URL)+" ("+h.Format+")")
|
||||
}
|
||||
fmt.Fprintf(&b, " alerts to %s\n", strings.Join(where, ", "))
|
||||
}
|
||||
granted := false
|
||||
for _, d := range m.db {
|
||||
addr := strings.TrimSpace(d.addr.Value())
|
||||
|
||||
@@ -45,7 +45,7 @@ func TestTheSetupAsksForWhatHighAvailabilityNeedsAndRefusesWhatWillNotDo(t *test
|
||||
t.Fatalf("high availability did not lead to the administrator: page %d", s.page)
|
||||
}
|
||||
// No email: refused, with a reason.
|
||||
m = press(m, "tab", "tab", "enter")
|
||||
m = press(m, "tab", "tab", "tab", "enter")
|
||||
if s := m.(setupModel); s.page != pAdmin || !strings.Contains(s.formErr, "email") {
|
||||
t.Fatalf("went on without an administrator: %+v", s.formErr)
|
||||
}
|
||||
@@ -54,7 +54,7 @@ func TestTheSetupAsksForWhatHighAvailabilityNeedsAndRefusesWhatWillNotDo(t *test
|
||||
if !strings.Contains(m.View(), "Not yet") {
|
||||
t.Error("the refusal was not shown as a dialog")
|
||||
}
|
||||
m = press(m, "x", "a@example.com", "tab", "tab", "enter")
|
||||
m = press(m, "x", "a@example.com", "tab", "tab", "tab", "enter")
|
||||
s := m.(setupModel)
|
||||
if s.page != pMachines {
|
||||
t.Fatalf("page %d: %s", s.page, s.formErr)
|
||||
|
||||
@@ -107,7 +107,7 @@ var Routes = []Route{
|
||||
{Method: "GET", Path: "/v1/memory/{key}", Summary: "A note", Response: MemoryNote{}},
|
||||
{Method: "GET", Path: "/v1/memory", Summary: "The project's notes", Response: MemoryList{}},
|
||||
{Method: "DELETE", Path: "/v1/memory/{key}", Summary: "Forget a note"},
|
||||
{Method: "POST", Path: "/mcp", Summary: "The gateway's MCP server: the memory notes, and pictures, 3D models, vision and workflows, as tools",
|
||||
{Method: "POST", Path: "/mcp", Summary: "The gateway's MCP server: the memory notes as tools",
|
||||
Request: MCPRequest{}, Response: anyJSON},
|
||||
{Method: "GET", Path: "/me", Summary: "The caller: its key, its place in the tree, its limits and what it has used",
|
||||
Response: Me{}},
|
||||
|
||||
@@ -258,7 +258,11 @@ type Result struct {
|
||||
|
||||
// redact keeps a webhook's host and drops its path, where Slack and others
|
||||
// put the secret.
|
||||
func redact(u string) string {
|
||||
func redact(u string) string { return Redact(u) }
|
||||
|
||||
// Redact is a webhook's address as it may be shown: its host, not the path
|
||||
// that is its secret.
|
||||
func Redact(u string) string {
|
||||
if i := strings.Index(u, "://"); i >= 0 {
|
||||
if j := strings.IndexByte(u[i+3:], '/'); j >= 0 {
|
||||
return u[:i+3+j] + "/…"
|
||||
|
||||
@@ -359,9 +359,7 @@ func (r *Router) Handler() http.Handler {
|
||||
}
|
||||
|
||||
if r.auth == nil {
|
||||
h := r.refuseWhileStopped(r.refuseLoops(mux))
|
||||
r.selfHandler.Store(&h)
|
||||
return h
|
||||
return r.refuseWhileStopped(r.refuseLoops(mux))
|
||||
}
|
||||
exempt := []string{"/health", "/ready", "/leader", "/v1/openapi.json", "/", "/about", "/_static/", "/_lang/", "/claim/",
|
||||
"/signup", "/login", "/login/", "/logout", "/profile", "/profile/", "/scim/",
|
||||
@@ -375,9 +373,7 @@ func (r *Router) Handler() http.Handler {
|
||||
// Wrapping the whole mux rather than checking inside each handler: adding a
|
||||
// route must not be able to add an unauthenticated one by omission. Health
|
||||
// stays open because a load balancer has no key.
|
||||
h := r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
|
||||
r.selfHandler.Store(&h)
|
||||
return h
|
||||
return r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
|
||||
}
|
||||
|
||||
// refuseWhileStopped answers every request for a model with 503 while the
|
||||
|
||||
@@ -1,937 +0,0 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
_ "image/jpeg" // a picture sent to look at may be a JPEG
|
||||
"image/png"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// --- MCP: pictures, 3D models, looking at pictures, workflows ---
|
||||
//
|
||||
// The gateway's own models offered as tools, so an agent — Claude Code
|
||||
// driving Blender, say — makes a concept, a mesh from it and has a render
|
||||
// looked at without a tool server of its own. Each tool makes its request
|
||||
// of the gateway as the caller would, under the caller's key: the same
|
||||
// allowlist, limits, licences left out, guardrails and accounting. What it
|
||||
// makes is kept as a file of the caller's, said by id and by the address
|
||||
// it is fetched from with the same key, since a tool server over HTTP
|
||||
// cannot write to the caller's disk.
|
||||
|
||||
const (
|
||||
// mcpBody bounds a tool call: up to four pictures to look at, as base64.
|
||||
mcpBody = 64 << 20
|
||||
// previewEdge is the long side of the picture shown to the agent beside
|
||||
// the file: enough to see it, not a megapixel of tokens.
|
||||
previewEdge = 512
|
||||
// lookEdge is the long side of a picture sent to a vision model:
|
||||
// 1280x720 is about 940 tokens.
|
||||
lookEdge = 1280
|
||||
// mcpWaitStreamed is how long a call answered as a stream waits for its
|
||||
// job, saying how far along it is meanwhile; mcpWaitPlain how long one
|
||||
// answered in one piece does, under the minute a proxy allows.
|
||||
mcpWaitStreamed = 10 * time.Minute
|
||||
mcpWaitPlain = 50 * time.Second
|
||||
mcpWaitMax = 15 * time.Minute
|
||||
mcpPoll = 2 * time.Second
|
||||
)
|
||||
|
||||
// The prompt endings of generate_image's styles: what a picture a mesh is
|
||||
// made from wants, and a surface that tiles.
|
||||
var imageStyles = map[string]string{
|
||||
"none": "",
|
||||
"concept": ", single object centered, orthographic three-quarter view, flat even studio lighting, no shadows, " +
|
||||
"no background elements, plain flat grey background, clean silhouette, game asset concept, high detail, no text, no watermark",
|
||||
"side": ", single object, exact side view (profile), orthographic, flat even studio lighting, no shadows, " +
|
||||
"plain flat grey background, clean silhouette, game asset concept, high detail, no text, no watermark",
|
||||
"texture": ", seamless tileable texture, top-down orthographic, flat even lighting, no shadows, no perspective, no text",
|
||||
}
|
||||
|
||||
// mcpMediaInstructions is said to an agent when it connects and the gateway
|
||||
// has something to make pictures or 3D models with.
|
||||
const mcpMediaInstructions = "Pictures, 3D models and a model that looks at pictures, on this gateway's own GPUs. " +
|
||||
"A typical loop: generate_image (style concept) -> generate_mesh from its file_id -> import the GLB -> render -> " +
|
||||
"review_images against the concept -> fix. Results are files: fetch them from their url with your API key. " +
|
||||
"AI meshes suit rocks, debris and background props; hard-surface hero assets are usually better built by hand " +
|
||||
"from the concept. A GPU holds one of these models at a time and switching costs a cold start, so batch calls " +
|
||||
"of the same kind; gateway_status says which is loaded."
|
||||
|
||||
// mediaModels are the gateway's models by what they do, by name, deduplicated
|
||||
// from the aliases that also resolve to them.
|
||||
type mediaModels struct{ images, meshes, vision, workflows []string }
|
||||
|
||||
func (r *Router) mediaModels() mediaModels {
|
||||
var out mediaModels
|
||||
seen := map[string]bool{}
|
||||
for _, m := range r.modelMap() {
|
||||
if seen[m.Name] || m.IsGuard() {
|
||||
continue
|
||||
}
|
||||
seen[m.Name] = true
|
||||
if m.MakesImages() {
|
||||
out.images = append(out.images, m.Name)
|
||||
}
|
||||
if m.MakesMeshes() {
|
||||
out.meshes = append(out.meshes, m.Name)
|
||||
}
|
||||
if m.Capabilities.Vision {
|
||||
out.vision = append(out.vision, m.Name)
|
||||
}
|
||||
if len(m.ComfyInputs()) > 0 {
|
||||
out.workflows = append(out.workflows, m.Name)
|
||||
}
|
||||
}
|
||||
for _, l := range [][]string{out.images, out.meshes, out.vision, out.workflows} {
|
||||
slices.Sort(l)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (mm mediaModels) any() bool {
|
||||
return len(mm.images)+len(mm.meshes)+len(mm.vision)+len(mm.workflows) > 0
|
||||
}
|
||||
|
||||
// defaultOf is the model a tool uses when the call names none: the one the
|
||||
// tool is best known with when the gateway has it, else the first.
|
||||
func defaultOf(models []string, preferred ...string) string {
|
||||
for _, p := range preferred {
|
||||
if slices.Contains(models, p) {
|
||||
return p
|
||||
}
|
||||
}
|
||||
if len(models) == 0 {
|
||||
return ""
|
||||
}
|
||||
return models[0]
|
||||
}
|
||||
|
||||
func enumOrNone(models []string) map[string]any {
|
||||
s := map[string]any{"type": "string"}
|
||||
if len(models) > 0 {
|
||||
s["enum"] = models
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// mediaTools are the tools for what the gateway's models can do: none for
|
||||
// what it has no model for.
|
||||
func (r *Router) mediaTools() []map[string]any {
|
||||
mm := r.mediaModels()
|
||||
var tools []map[string]any
|
||||
fileArg := map[string]any{"type": "string",
|
||||
"description": "A picture: a /v1/files id (file-…), a data URL, or base64"}
|
||||
wait := map[string]any{"type": "integer", "minimum": 0, "maximum": int(mcpWaitMax.Seconds()),
|
||||
"description": "Most seconds to wait for it before answering with the job to ask after (job_status)"}
|
||||
if len(mm.images) > 0 {
|
||||
tools = append(tools, map[string]any{
|
||||
"name": "generate_image",
|
||||
"description": "Make a picture from a prompt (~15 s warm, longer when the GPU must switch models). " +
|
||||
"Keeps the PNG as a file and answers its file_id, url and seed, with a small preview. " +
|
||||
"style: 'concept' (one object, 3/4 view, grey background: good for generate_mesh), 'side' " +
|
||||
"(exact side view: best meshes for vehicles), 'texture' (seamless tileable surface), 'none' (the prompt as is).",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"prompt"},
|
||||
"properties": map[string]any{
|
||||
"prompt": map[string]any{"type": "string"},
|
||||
"style": map[string]any{"type": "string", "enum": []string{"none", "concept", "side", "texture"}, "default": "none"},
|
||||
"size": map[string]any{"type": "integer", "description": "Square side in px, a multiple of 16, 256-2048", "default": 1024},
|
||||
"seed": map[string]any{"type": "integer"},
|
||||
"negative_prompt": map[string]any{"type": "string"},
|
||||
"model": enumOrNone(mm.images),
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
if len(mm.meshes) > 0 {
|
||||
tools = append(tools, map[string]any{
|
||||
"name": "generate_mesh",
|
||||
"description": "Make a 3D model (GLB) from a picture of one object on a plain background — the gateway " +
|
||||
"cuts the background out itself — or from a prompt, for which a concept picture is made first. " +
|
||||
"Takes 30 s to a few minutes; waits for it, saying how far along it is when the client takes a stream. " +
|
||||
"Answers the GLB's file_id and url, and its licence. The mesh is about 1 m across, its orientation arbitrary.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"image": fileArg,
|
||||
"prompt": map[string]any{"type": "string", "description": "Used when no image is given"},
|
||||
"model": enumOrNone(mm.meshes),
|
||||
"seed": map[string]any{"type": "integer"},
|
||||
"resolution": map[string]any{"type": "integer", "description": "The model's own grid, a multiple of 32 (TripoSR: 256 is enough)"},
|
||||
"texture_resolution": map[string]any{"type": "integer", "minimum": 512, "maximum": 2048,
|
||||
"description": "For a model that textures its mesh (sf3d): the texture's side in px"},
|
||||
"inputs": map[string]any{"type": "object",
|
||||
"description": "Any other input the model's workflow declares (gateway_status lists them): remesh, vertex_count, foreground_ratio, albedo_match…"},
|
||||
"wait_seconds": wait,
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
if len(mm.vision) > 0 {
|
||||
tools = append(tools, map[string]any{
|
||||
"name": "review_images",
|
||||
"description": "Ask a vision model about 1-4 pictures — a render beside its concept, say. A cheap critic " +
|
||||
"for render-and-fix loops. json_schema makes it answer strict JSON. A cold start can take minutes " +
|
||||
"when the GPU was making pictures.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"images", "question"},
|
||||
"properties": map[string]any{
|
||||
"images": map[string]any{"type": "array", "items": fileArg, "minItems": 1, "maxItems": 4},
|
||||
"question": map[string]any{"type": "string"},
|
||||
"json_schema": map[string]any{"type": "object", "description": "JSON Schema the answer must follow"},
|
||||
"model": enumOrNone(mm.vision),
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
if len(mm.workflows) > 0 {
|
||||
tools = append(tools, map[string]any{
|
||||
"name": "run_workflow",
|
||||
"description": "Run a ComfyUI workflow the gateway serves, with the inputs its catalog entry declares " +
|
||||
"(gateway_status lists them). A file input is a /v1/files id, a data URL or base64. Answers its result file.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"model", "inputs"},
|
||||
"properties": map[string]any{
|
||||
"model": enumOrNone(mm.workflows),
|
||||
"inputs": map[string]any{"type": "object"},
|
||||
"wait_seconds": wait,
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
if mm.any() {
|
||||
tools = append(tools,
|
||||
map[string]any{
|
||||
"name": "job_status",
|
||||
"description": "Where a 3D model or workflow run stands, and its result file when done. wait_seconds waits for it to finish.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"job_id"},
|
||||
"properties": map[string]any{"job_id": map[string]any{"type": "string"}, "wait_seconds": wait},
|
||||
},
|
||||
},
|
||||
map[string]any{
|
||||
"name": "gateway_status",
|
||||
"description": "The gateway's picture, 3D, vision and workflow models: whether each is loaded on a GPU now, " +
|
||||
"and a workflow's inputs; and your last few jobs. Use it to batch calls and avoid cold starts.",
|
||||
"inputSchema": map[string]any{"type": "object", "properties": map[string]any{}},
|
||||
})
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// mediaCall is one tool call: who made it, and how to tell them how far it
|
||||
// has got.
|
||||
type mediaCall struct {
|
||||
req *http.Request
|
||||
progress func(percent float64, message string)
|
||||
// streamed is a call whose answer is a stream, which may wait longer.
|
||||
streamed bool
|
||||
}
|
||||
|
||||
func (c mediaCall) wait(asked *int) time.Duration {
|
||||
d := mcpWaitPlain
|
||||
if c.streamed {
|
||||
d = mcpWaitStreamed
|
||||
}
|
||||
if asked != nil {
|
||||
d = min(time.Duration(*asked)*time.Second, mcpWaitMax)
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// isMediaTool is whether a tool is one of these.
|
||||
func isMediaTool(name string) bool {
|
||||
switch name {
|
||||
case "generate_image", "generate_mesh", "review_images", "run_workflow", "job_status", "gateway_status":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// callMediaTool runs one of the tools. As with the notes, a failure is a
|
||||
// tool result with isError set, which the model gets to read.
|
||||
func (r *Router) callMediaTool(c mediaCall, name string, args json.RawMessage) map[string]any {
|
||||
var res map[string]any
|
||||
var err error
|
||||
switch name {
|
||||
case "generate_image":
|
||||
res, err = r.toolGenerateImage(c, args)
|
||||
case "generate_mesh":
|
||||
res, err = r.toolGenerateMesh(c, args)
|
||||
case "review_images":
|
||||
res, err = r.toolReviewImages(c, args)
|
||||
case "run_workflow":
|
||||
res, err = r.toolRunWorkflow(c, args)
|
||||
case "job_status":
|
||||
res, err = r.toolJobStatus(c, args)
|
||||
case "gateway_status":
|
||||
res, err = r.toolGatewayStatus(c)
|
||||
default:
|
||||
err = errors.New("unknown tool " + name)
|
||||
}
|
||||
if err != nil {
|
||||
return toolText(err.Error(), true)
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// self makes a request of the gateway's client API as the caller: their key,
|
||||
// their agent name, and the context the call was made in.
|
||||
func (r *Router) self(c mediaCall, method, path string, body any) (int, http.Header, []byte, error) {
|
||||
h := r.selfHandler.Load()
|
||||
if h == nil {
|
||||
return 0, nil, nil, errors.New("the gateway's API is not mounted here")
|
||||
}
|
||||
var rd *bytes.Reader
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return 0, nil, nil, err
|
||||
}
|
||||
rd = bytes.NewReader(b)
|
||||
} else {
|
||||
rd = bytes.NewReader(nil)
|
||||
}
|
||||
out, err := http.NewRequestWithContext(c.req.Context(), method, path, rd)
|
||||
if err != nil {
|
||||
return 0, nil, nil, err
|
||||
}
|
||||
for _, k := range []string{"Authorization", "X-Api-Key", agentHeader, "X-Forwarded-For", "X-Real-Ip"} {
|
||||
if v := c.req.Header.Get(k); v != "" {
|
||||
out.Header.Set(k, v)
|
||||
}
|
||||
}
|
||||
out.RemoteAddr = c.req.RemoteAddr
|
||||
if body != nil {
|
||||
out.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
// A picture or a mesh comes back whole: past the 1 MiB a judge's
|
||||
// answer is held to.
|
||||
w := &bufferWriter{header: http.Header{}, max: mcpBody * 2}
|
||||
(*h).ServeHTTP(w, out)
|
||||
status := w.status
|
||||
if status == 0 {
|
||||
status = http.StatusOK
|
||||
}
|
||||
return status, w.header, w.buf.Bytes(), nil
|
||||
}
|
||||
|
||||
// selfJSON is self for an answer in JSON, a failure said as the API said it.
|
||||
func (r *Router) selfJSON(c mediaCall, method, path string, body, into any) (http.Header, error) {
|
||||
status, h, b, err := r.self(c, method, path, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status >= 300 {
|
||||
return h, fmt.Errorf("%s %s: %d %s", method, path, status, errorMessage(b))
|
||||
}
|
||||
if into == nil {
|
||||
return h, nil
|
||||
}
|
||||
return h, json.Unmarshal(b, into)
|
||||
}
|
||||
|
||||
// fileURL is where a file is fetched from, with the caller's key, as the
|
||||
// caller reached the gateway.
|
||||
func fileURL(req *http.Request, id string) string {
|
||||
scheme := "http"
|
||||
if req.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
if p := req.Header.Get("X-Forwarded-Proto"); p != "" {
|
||||
scheme = strings.TrimSpace(strings.Split(p, ",")[0])
|
||||
}
|
||||
host := req.Host
|
||||
if h := req.Header.Get("X-Forwarded-Host"); h != "" {
|
||||
host = strings.TrimSpace(strings.Split(h, ",")[0])
|
||||
}
|
||||
return scheme + "://" + host + "/v1/files/" + id + "/content"
|
||||
}
|
||||
|
||||
// keepFile keeps what a tool made as the caller's file, for as long as a
|
||||
// job's result is kept.
|
||||
func (r *Router) keepFile(c mediaCall, name string, content []byte) (store.File, error) {
|
||||
p, ok := auth.FromContext(c.req.Context())
|
||||
if !ok || r.store == nil {
|
||||
return store.File{}, errors.New("this gateway keeps no files (it runs without authentication or a store)")
|
||||
}
|
||||
if r.noStore(callerSubject(c.req)) {
|
||||
return store.File{}, errors.New("the gateway's policy for this caller keeps nothing, so no file is kept")
|
||||
}
|
||||
_, _, _, days := r.batchCfg()
|
||||
now := time.Now()
|
||||
f := store.File{ID: newID("file-"), Owner: ownerOf(p), Purpose: purposeJobOutput, Filename: name,
|
||||
Bytes: int64(len(content)), CreatedAt: now, ExpiresAt: now.AddDate(0, 0, days)}
|
||||
return f, r.store.PutFile(c.req.Context(), f, content)
|
||||
}
|
||||
|
||||
// picture reads a picture argument: a file of the caller's by id, a data URL,
|
||||
// or base64.
|
||||
func (r *Router) picture(c mediaCall, s string) ([]byte, error) {
|
||||
s = strings.TrimSpace(s)
|
||||
switch {
|
||||
case s == "":
|
||||
return nil, errors.New("no picture given")
|
||||
case strings.HasPrefix(s, "file-"):
|
||||
p, ok := auth.FromContext(c.req.Context())
|
||||
if !ok || r.store == nil {
|
||||
return nil, errors.New("this gateway keeps no files: send the picture as base64 or a data URL")
|
||||
}
|
||||
f, err := r.store.FileByID(c.req.Context(), s)
|
||||
if errors.Is(err, store.ErrNotFound) || err == nil && f.Owner != ownerOf(p) {
|
||||
return nil, fmt.Errorf("no file %s of yours", s)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.store.FileContent(c.req.Context(), s)
|
||||
case strings.HasPrefix(s, "data:"):
|
||||
_, after, ok := strings.Cut(s, ";base64,")
|
||||
if !ok {
|
||||
return nil, errors.New("a data URL without base64 in it")
|
||||
}
|
||||
s = after
|
||||
}
|
||||
b, err := base64.StdEncoding.DecodeString(s)
|
||||
if err != nil {
|
||||
return nil, errors.New("the picture is neither a file id, a data URL nor base64")
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func dataURL(png []byte) string {
|
||||
return "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)
|
||||
}
|
||||
|
||||
// shrink returns a picture as a PNG whose long side is at most edge.
|
||||
func shrink(pic []byte, edge int) ([]byte, error) {
|
||||
src, format, err := image.Decode(bytes.NewReader(pic))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("not a picture this gateway reads (PNG or JPEG): %w", err)
|
||||
}
|
||||
b := src.Bounds()
|
||||
w, h := b.Dx(), b.Dy()
|
||||
if max(w, h) <= edge {
|
||||
if format == "png" {
|
||||
return pic, nil
|
||||
}
|
||||
edge = max(w, h) // only re-encoded
|
||||
}
|
||||
scale := float64(edge) / float64(max(w, h))
|
||||
dw, dh := max(1, int(float64(w)*scale+0.5)), max(1, int(float64(h)*scale+0.5))
|
||||
dst := image.NewNRGBA(image.Rect(0, 0, dw, dh))
|
||||
// Each pixel the average of the box of pixels it stands for.
|
||||
for y := range dh {
|
||||
y0, y1 := b.Min.Y+y*h/dh, b.Min.Y+max((y+1)*h/dh, y*h/dh+1)
|
||||
for x := range dw {
|
||||
x0, x1 := b.Min.X+x*w/dw, b.Min.X+max((x+1)*w/dw, x*w/dw+1)
|
||||
var sr, sg, sb, sa, n uint64
|
||||
for yy := y0; yy < y1; yy++ {
|
||||
for xx := x0; xx < x1; xx++ {
|
||||
cr, cg, cb, ca := src.At(xx, yy).RGBA()
|
||||
sr, sg, sb, sa, n = sr+uint64(cr), sg+uint64(cg), sb+uint64(cb), sa+uint64(ca), n+1
|
||||
}
|
||||
}
|
||||
// Premultiplied sums, back to straight colour.
|
||||
if sa == 0 {
|
||||
continue
|
||||
}
|
||||
dst.SetNRGBA(x, y, color.NRGBA{
|
||||
R: uint8(sr * 0xff / sa), G: uint8(sg * 0xff / sa), B: uint8(sb * 0xff / sa), A: uint8(sa / n >> 8),
|
||||
})
|
||||
}
|
||||
}
|
||||
var out bytes.Buffer
|
||||
if err := png.Encode(&out, dst); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out.Bytes(), nil
|
||||
}
|
||||
|
||||
// --- the tools ---
|
||||
|
||||
func (r *Router) toolGenerateImage(c mediaCall, raw json.RawMessage) (map[string]any, error) {
|
||||
var a struct {
|
||||
Prompt string `json:"prompt"`
|
||||
Style string `json:"style"`
|
||||
NegativePrompt string `json:"negative_prompt"`
|
||||
Model string `json:"model"`
|
||||
Size int `json:"size"`
|
||||
Seed *int64 `json:"seed"`
|
||||
}
|
||||
if err := decodeArgs(raw, &a); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
res, _, err := r.makeImage(c, a.Prompt, a.Style, a.NegativePrompt, a.Model, a.Size, a.Seed)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// makeImage makes a picture and keeps it: the tool's answer, and the PNG.
|
||||
func (r *Router) makeImage(c mediaCall, prompt, style, negative, model string, size int, seed *int64) (map[string]any, []byte, error) {
|
||||
if strings.TrimSpace(prompt) == "" {
|
||||
return nil, nil, errors.New("prompt is required")
|
||||
}
|
||||
if style == "" {
|
||||
style = "none"
|
||||
}
|
||||
suffix, ok := imageStyles[style]
|
||||
if !ok {
|
||||
return nil, nil, fmt.Errorf("style %q: concept, side, texture or none", style)
|
||||
}
|
||||
if model == "" {
|
||||
model = defaultOf(r.mediaModels().images, "flux-schnell")
|
||||
}
|
||||
if size == 0 {
|
||||
size = 1024
|
||||
}
|
||||
body := map[string]any{"model": model, "prompt": prompt + suffix, "n": 1,
|
||||
"size": fmt.Sprintf("%dx%d", size, size), "response_format": "b64_json"}
|
||||
if seed != nil {
|
||||
body["seed"] = *seed
|
||||
}
|
||||
if negative != "" {
|
||||
body["negative_prompt"] = negative
|
||||
}
|
||||
c.progress(5, "making the picture with "+model)
|
||||
var out struct {
|
||||
Data []struct {
|
||||
B64 string `json:"b64_json"`
|
||||
} `json:"data"`
|
||||
}
|
||||
h, err := r.selfJSON(c, http.MethodPost, imagesPath, body, &out)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if len(out.Data) == 0 {
|
||||
return nil, nil, errors.New(model + " made no picture")
|
||||
}
|
||||
pic, err := base64.StdEncoding.DecodeString(out.Data[0].B64)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
info := map[string]any{"model": model, "width": size, "height": size}
|
||||
if s, err := strconv.ParseInt(h.Get("x-llmgw-seed"), 10, 64); err == nil {
|
||||
info["seed"] = s
|
||||
}
|
||||
if m, ok := r.modelMap()[model]; ok && m.License != nil {
|
||||
info["license"] = m.License
|
||||
}
|
||||
content := []map[string]any{}
|
||||
if f, err := r.keepFile(c, "image-"+strings.TrimPrefix(newID(""), "-")+".png", pic); err == nil {
|
||||
info["file_id"], info["url"], info["bytes"] = f.ID, fileURL(c.req, f.ID), f.Bytes
|
||||
} else {
|
||||
// Nowhere to keep it: the picture itself, whole.
|
||||
info["note"] = err.Error() + "; the picture is in this answer at full size"
|
||||
content = append(content, map[string]any{"type": "image", "mimeType": "image/png", "data": out.Data[0].B64})
|
||||
}
|
||||
text, _ := json.Marshal(info)
|
||||
content = append([]map[string]any{{"type": "text", "text": string(text)}}, content...)
|
||||
if _, kept := info["file_id"]; kept {
|
||||
if small, err := shrink(pic, previewEdge); err == nil {
|
||||
content = append(content, map[string]any{"type": "image", "mimeType": "image/png",
|
||||
"data": base64.StdEncoding.EncodeToString(small)})
|
||||
}
|
||||
}
|
||||
return map[string]any{"content": content, "isError": false}, pic, nil
|
||||
}
|
||||
|
||||
func (r *Router) toolGenerateMesh(c mediaCall, raw json.RawMessage) (map[string]any, error) {
|
||||
var a struct {
|
||||
Image, Prompt, Model string
|
||||
Seed *int64
|
||||
Resolution int
|
||||
TextureResolution int `json:"texture_resolution"`
|
||||
Inputs map[string]any
|
||||
WaitSeconds *int `json:"wait_seconds"`
|
||||
}
|
||||
if err := decodeArgs(raw, &a); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
mm := r.mediaModels()
|
||||
if a.Model == "" {
|
||||
a.Model = defaultOf(mm.meshes, "sf3d", "triposr")
|
||||
}
|
||||
m, ok := r.modelMap()[a.Model]
|
||||
if !ok || !m.MakesMeshes() {
|
||||
return nil, fmt.Errorf("%q makes no 3D models here; these do: %s", a.Model, strings.Join(mm.meshes, ", "))
|
||||
}
|
||||
extra := map[string]any{}
|
||||
var pic []byte
|
||||
var err error
|
||||
switch {
|
||||
case a.Image != "":
|
||||
if pic, err = r.picture(c, a.Image); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case a.Prompt != "":
|
||||
res, made, err := r.makeImage(c, a.Prompt, "concept", "", "", 0, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("the concept picture: %w", err)
|
||||
}
|
||||
pic = made
|
||||
var info map[string]any
|
||||
if json.Unmarshal([]byte(res["content"].([]map[string]any)[0]["text"].(string)), &info) == nil {
|
||||
extra["concept_file_id"], extra["concept_seed"] = info["file_id"], info["seed"]
|
||||
}
|
||||
default:
|
||||
return nil, errors.New("an image, or a prompt to make one from")
|
||||
}
|
||||
|
||||
// A model whose workflow declares its inputs takes them at the
|
||||
// workflows' path, every one it names; one that does not, the 3D path's
|
||||
// seed and grid.
|
||||
path, body := meshesPath, map[string]any{"model": a.Model, "image": dataURL(pic)}
|
||||
if declared := m.ComfyInputs(); len(declared) > 0 {
|
||||
names := map[string]string{}
|
||||
var image string
|
||||
for _, in := range declared {
|
||||
names[in.Name] = in.Type
|
||||
if image == "" && in.Type == "image" {
|
||||
image = in.Name
|
||||
}
|
||||
}
|
||||
if image == "" {
|
||||
return nil, fmt.Errorf("%s's workflow takes no picture", a.Model)
|
||||
}
|
||||
inputs := map[string]any{image: dataURL(pic)}
|
||||
for k, v := range a.Inputs {
|
||||
inputs[k] = v
|
||||
}
|
||||
|
||||
// A textured mesh's resolution is its texture's; a grid's, the grid's.
|
||||
if _, ok := names["resolution"]; ok {
|
||||
switch {
|
||||
case a.TextureResolution > 0:
|
||||
inputs["resolution"] = a.TextureResolution
|
||||
case a.Resolution > 0:
|
||||
inputs["resolution"] = a.Resolution
|
||||
}
|
||||
}
|
||||
path, body = workflowsPath, map[string]any{"model": a.Model, "inputs": inputs}
|
||||
if a.Seed != nil {
|
||||
body["seed"] = *a.Seed
|
||||
}
|
||||
} else {
|
||||
if a.Seed != nil {
|
||||
body["seed"] = *a.Seed
|
||||
}
|
||||
if a.Resolution > 0 {
|
||||
body["resolution"] = a.Resolution
|
||||
}
|
||||
}
|
||||
var job struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if _, err := r.selfJSON(c, http.MethodPost, path, body, &job); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.progress(10, "queued as "+job.ID)
|
||||
return r.waitJob(c, job.ID, c.wait(a.WaitSeconds), extra)
|
||||
}
|
||||
|
||||
func (r *Router) toolRunWorkflow(c mediaCall, raw json.RawMessage) (map[string]any, error) {
|
||||
var a struct {
|
||||
Model string
|
||||
Inputs map[string]any
|
||||
WaitSeconds *int `json:"wait_seconds"`
|
||||
}
|
||||
if err := decodeArgs(raw, &a); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m, ok := r.modelMap()[a.Model]
|
||||
if !ok || len(m.ComfyInputs()) == 0 {
|
||||
return nil, fmt.Errorf("%q runs no workflow here; these do: %s", a.Model, strings.Join(r.mediaModels().workflows, ", "))
|
||||
}
|
||||
// A file input given as a file id goes as it is: the API takes those.
|
||||
var job struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
if _, err := r.selfJSON(c, http.MethodPost, workflowsPath, map[string]any{"model": a.Model, "inputs": a.Inputs}, &job); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.progress(10, "queued as "+job.ID)
|
||||
return r.waitJob(c, job.ID, c.wait(a.WaitSeconds), nil)
|
||||
}
|
||||
|
||||
func (r *Router) toolJobStatus(c mediaCall, raw json.RawMessage) (map[string]any, error) {
|
||||
var a struct {
|
||||
JobID string `json:"job_id"`
|
||||
WaitSeconds *int `json:"wait_seconds"`
|
||||
}
|
||||
if err := decodeArgs(raw, &a); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if a.JobID == "" {
|
||||
return nil, errors.New("job_id is required")
|
||||
}
|
||||
d := time.Duration(0)
|
||||
if a.WaitSeconds != nil {
|
||||
d = c.wait(a.WaitSeconds)
|
||||
}
|
||||
return r.waitJob(c, a.JobID, d, nil)
|
||||
}
|
||||
|
||||
// jobView is the part of a job a tool says.
|
||||
type jobView struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Status string `json:"status"`
|
||||
Progress json.RawMessage `json:"progress"`
|
||||
OutputFileID *string `json:"output_file_id"`
|
||||
ContentType *string `json:"content_type"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
License json.RawMessage `json:"license"`
|
||||
}
|
||||
|
||||
// waitJob asks after a job until it is done or d has passed, saying how far
|
||||
// along it is, and answers where it stands: its file when done.
|
||||
func (r *Router) waitJob(c mediaCall, id string, d time.Duration, extra map[string]any) (map[string]any, error) {
|
||||
deadline := time.Now().Add(d)
|
||||
var j jobView
|
||||
for {
|
||||
if _, err := r.selfJSON(c, http.MethodGet, workflowsPath+"/"+id, nil, &j); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if store.JobDone(j.Status) || !time.Now().Before(deadline) {
|
||||
break
|
||||
}
|
||||
pct, msg := jobProgress(j)
|
||||
c.progress(pct, msg)
|
||||
select {
|
||||
case <-c.req.Context().Done():
|
||||
return nil, c.req.Context().Err()
|
||||
case <-time.After(mcpPoll):
|
||||
}
|
||||
}
|
||||
info := map[string]any{"job_id": j.ID, "model": j.Model, "status": j.Status}
|
||||
for k, v := range extra {
|
||||
if v != nil {
|
||||
info[k] = v
|
||||
}
|
||||
}
|
||||
if len(j.License) > 0 && string(j.License) != "null" {
|
||||
info["license"] = j.License
|
||||
}
|
||||
switch {
|
||||
case j.Status == store.JobCompleted && j.OutputFileID != nil:
|
||||
info["file_id"], info["url"] = *j.OutputFileID, fileURL(c.req, *j.OutputFileID)
|
||||
if j.ContentType != nil {
|
||||
info["content_type"] = *j.ContentType
|
||||
}
|
||||
if r.store != nil {
|
||||
if f, err := r.store.FileByID(c.req.Context(), *j.OutputFileID); err == nil {
|
||||
info["bytes"] = f.Bytes
|
||||
}
|
||||
}
|
||||
case store.JobDone(j.Status):
|
||||
if j.Error != nil {
|
||||
info["error"] = j.Error.Message
|
||||
}
|
||||
b, _ := json.Marshal(info)
|
||||
return toolText(string(b), true), nil
|
||||
default:
|
||||
info["progress"] = j.Progress
|
||||
info["next"] = "not done yet: call job_status with this job_id (and wait_seconds) to wait for it"
|
||||
}
|
||||
b, _ := json.Marshal(info)
|
||||
return toolText(string(b), false), nil
|
||||
}
|
||||
|
||||
// jobProgress reads a job's progress as a percent and a phrase.
|
||||
func jobProgress(j jobView) (float64, string) {
|
||||
var p struct {
|
||||
State string `json:"state"`
|
||||
Ahead int `json:"ahead"`
|
||||
Percent float64 `json:"percent"`
|
||||
Node string `json:"node_class"`
|
||||
HeldBy string `json:"held_by"`
|
||||
}
|
||||
_ = json.Unmarshal(j.Progress, &p)
|
||||
switch {
|
||||
case p.State == "queued":
|
||||
return 10, fmt.Sprintf("%s: queued, %d ahead", j.ID, p.Ahead)
|
||||
case p.State == "waiting_gpu":
|
||||
return 10, j.ID + ": waiting for the GPU, held by " + p.HeldBy
|
||||
case p.State == "running":
|
||||
return 10 + 0.9*p.Percent, fmt.Sprintf("%s: %s %.0f%%", j.ID, p.Node, p.Percent)
|
||||
}
|
||||
return 10, j.ID + ": " + j.Status
|
||||
}
|
||||
|
||||
func (r *Router) toolReviewImages(c mediaCall, raw json.RawMessage) (map[string]any, error) {
|
||||
var a struct {
|
||||
Images []string
|
||||
Question string
|
||||
JSONSchema map[string]any `json:"json_schema"`
|
||||
Model string
|
||||
}
|
||||
if err := decodeArgs(raw, &a); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(a.Images) == 0 || len(a.Images) > 4 {
|
||||
return nil, errors.New("one to four images")
|
||||
}
|
||||
if strings.TrimSpace(a.Question) == "" {
|
||||
return nil, errors.New("question is required")
|
||||
}
|
||||
if a.Model == "" {
|
||||
a.Model = defaultOf(r.mediaModels().vision, "vlm")
|
||||
}
|
||||
if a.Model == "" {
|
||||
return nil, errors.New("no model here looks at pictures")
|
||||
}
|
||||
parts := []map[string]any{{"type": "text", "text": a.Question}}
|
||||
for i, s := range a.Images {
|
||||
pic, err := r.picture(c, s)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("image %d: %w", i+1, err)
|
||||
}
|
||||
small, err := shrink(pic, lookEdge)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("image %d: %w", i+1, err)
|
||||
}
|
||||
parts = append(parts, map[string]any{"type": "image_url", "image_url": map[string]any{"url": dataURL(small)}})
|
||||
}
|
||||
body := map[string]any{"model": a.Model, "messages": []map[string]any{{"role": "user", "content": parts}}}
|
||||
if a.JSONSchema != nil {
|
||||
body["response_format"] = map[string]any{"type": "json_schema",
|
||||
"json_schema": map[string]any{"name": "answer", "schema": a.JSONSchema, "strict": true}}
|
||||
}
|
||||
c.progress(10, "asking "+a.Model)
|
||||
var out struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if _, err := r.selfJSON(c, http.MethodPost, "/v1/chat/completions", body, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(out.Choices) == 0 {
|
||||
return nil, errors.New(a.Model + " gave no answer")
|
||||
}
|
||||
return toolText(out.Choices[0].Message.Content, false), nil
|
||||
}
|
||||
|
||||
func (r *Router) toolGatewayStatus(c mediaCall) (map[string]any, error) {
|
||||
var models struct {
|
||||
Data []map[string]any `json:"data"`
|
||||
}
|
||||
if _, err := r.selfJSON(c, http.MethodGet, "/v1/models", nil, &models); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var listed []map[string]any
|
||||
for _, m := range models.Data {
|
||||
caps, _ := m["capabilities"].(map[string]any)
|
||||
if caps["images"] != true && caps["meshes"] != true && caps["vision"] != true && m["workflow"] == nil {
|
||||
continue
|
||||
}
|
||||
e := map[string]any{"id": m["id"], "loaded": m["loaded"], "capabilities": caps}
|
||||
if w := m["workflow"]; w != nil {
|
||||
e["workflow"] = w
|
||||
}
|
||||
if l := m["license"]; l != nil {
|
||||
e["license"] = l
|
||||
}
|
||||
listed = append(listed, e)
|
||||
}
|
||||
out := map[string]any{"models": listed}
|
||||
var jobs struct {
|
||||
Data []jobView `json:"data"`
|
||||
}
|
||||
if _, err := r.selfJSON(c, http.MethodGet, workflowsPath+"?limit=5", nil, &jobs); err == nil {
|
||||
recent := make([]map[string]any, 0, len(jobs.Data))
|
||||
for _, j := range jobs.Data {
|
||||
recent = append(recent, map[string]any{"id": j.ID, "model": j.Model, "status": j.Status})
|
||||
}
|
||||
out["recent_jobs"] = recent
|
||||
}
|
||||
return toolJSON(out), nil
|
||||
}
|
||||
|
||||
// decodeArgs reads a tool's arguments, by their JSON names: a field's own
|
||||
// name in lower case unless its tag says.
|
||||
func decodeArgs(raw json.RawMessage, into any) error {
|
||||
if len(raw) == 0 || string(raw) == "null" {
|
||||
raw = []byte("{}")
|
||||
}
|
||||
if err := json.Unmarshal(raw, into); err != nil {
|
||||
return fmt.Errorf("the arguments: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// streamWriter answers a tool call as server-sent events: progress
|
||||
// notifications while it runs, then the answer.
|
||||
type streamWriter struct {
|
||||
mu sync.Mutex
|
||||
w http.ResponseWriter
|
||||
token json.RawMessage
|
||||
}
|
||||
|
||||
func (s *streamWriter) write(text string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
_, _ = fmt.Fprint(s.w, text)
|
||||
if f, ok := s.w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *streamWriter) event(v any) {
|
||||
b, _ := json.Marshal(v)
|
||||
s.write("event: message\ndata: " + string(b) + "\n\n")
|
||||
}
|
||||
|
||||
// keepAlive writes a comment every while until stop is closed: a picture
|
||||
// whose model starts cold is minutes of nothing said, which a proxy reads as
|
||||
// a connection gone idle.
|
||||
func (s *streamWriter) keepAlive(stop <-chan struct{}, every time.Duration) {
|
||||
t := time.NewTicker(every)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-t.C:
|
||||
s.write(": still working\n\n")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *streamWriter) progress(percent float64, message string) {
|
||||
if len(s.token) == 0 {
|
||||
return
|
||||
}
|
||||
s.event(map[string]any{"jsonrpc": "2.0", "method": "notifications/progress", "params": map[string]any{
|
||||
"progressToken": s.token, "progress": percent, "total": 100, "message": message}})
|
||||
}
|
||||
|
||||
// acceptsStream is a client that takes a POST's answer as a stream, as the
|
||||
// streamable HTTP transport lets it.
|
||||
func acceptsStream(req *http.Request) bool {
|
||||
return strings.Contains(req.Header.Get("Accept"), "text/event-stream")
|
||||
}
|
||||
@@ -1,266 +0,0 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/png"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
)
|
||||
|
||||
// mcp makes one JSON-RPC call at /mcp and answers its result; streamed, as a
|
||||
// client that takes server-sent events, it answers the events' data too.
|
||||
func mcp(t *testing.T, h http.Handler, secret, method string, params any, streamed bool) (map[string]any, []string) {
|
||||
t.Helper()
|
||||
b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": method, "params": params})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(b))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if secret != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+secret)
|
||||
}
|
||||
if streamed {
|
||||
req.Header.Set("Accept", "application/json, text/event-stream")
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("%s: %d %s", method, rec.Code, rec.Body.String())
|
||||
}
|
||||
body := rec.Body.String()
|
||||
var events []string
|
||||
if streamed {
|
||||
if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" {
|
||||
t.Fatalf("a streamed call answered %s: %s", ct, body)
|
||||
}
|
||||
for _, line := range strings.Split(body, "\n") {
|
||||
if d, ok := strings.CutPrefix(line, "data: "); ok {
|
||||
events = append(events, d)
|
||||
}
|
||||
}
|
||||
if len(events) == 0 {
|
||||
t.Fatalf("no events: %s", body)
|
||||
}
|
||||
body = events[len(events)-1]
|
||||
}
|
||||
var resp struct {
|
||||
Result map[string]any `json:"result"`
|
||||
Error any `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(body), &resp); err != nil || resp.Error != nil {
|
||||
t.Fatalf("%s: %v %s", method, err, body)
|
||||
}
|
||||
return resp.Result, events
|
||||
}
|
||||
|
||||
func toolNames(res map[string]any) []string {
|
||||
var names []string
|
||||
for _, tl := range res["tools"].([]any) {
|
||||
names = append(names, tl.(map[string]any)["name"].(string))
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// resultInfo is a tool result's first text, as the JSON it holds.
|
||||
func resultInfo(t *testing.T, res map[string]any) map[string]any {
|
||||
t.Helper()
|
||||
if res["isError"] == true {
|
||||
t.Fatalf("the tool failed: %v", res["content"])
|
||||
}
|
||||
var info map[string]any
|
||||
text := res["content"].([]any)[0].(map[string]any)["text"].(string)
|
||||
if err := json.Unmarshal([]byte(text), &info); err != nil {
|
||||
t.Fatalf("not JSON: %s", text)
|
||||
}
|
||||
return info
|
||||
}
|
||||
|
||||
func TestTheMCPToolsAreWhatTheGatewaysModelsCanDo(t *testing.T) {
|
||||
c := newMeshComfy(t)
|
||||
rt, secret := meshRouter(t, c)
|
||||
h := rt.Handler()
|
||||
init, _ := mcp(t, h, secret, "initialize", map[string]any{"protocolVersion": "2025-06-18"}, false)
|
||||
if init["protocolVersion"] != "2025-06-18" || !strings.Contains(init["instructions"].(string), "batch calls") {
|
||||
t.Errorf("initialize: %v", init)
|
||||
}
|
||||
list, _ := mcp(t, h, secret, "tools/list", map[string]any{}, false)
|
||||
names := strings.Join(toolNames(list), ",")
|
||||
for _, want := range []string{"remember", "generate_image", "generate_mesh", "job_status", "gateway_status"} {
|
||||
if !strings.Contains(names, want) {
|
||||
t.Errorf("no %s in %s", want, names)
|
||||
}
|
||||
}
|
||||
// Nothing here looks at pictures or declares a workflow's inputs.
|
||||
if strings.Contains(names, "review_images") || strings.Contains(names, "run_workflow") {
|
||||
t.Errorf("tools for models the gateway has not: %s", names)
|
||||
}
|
||||
|
||||
// A gateway with only text models offers the notes alone.
|
||||
plain := newTestRouter(t, baseConfig()).Handler()
|
||||
list, _ = mcp(t, plain, "", "tools/list", map[string]any{}, false)
|
||||
if got := strings.Join(toolNames(list), ","); got != "remember,recall,forget" {
|
||||
t.Errorf("a text-only gateway offers %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMeshIsMadeThroughMCPAndKeptAsTheCallersFile(t *testing.T) {
|
||||
c := newMeshComfy(t)
|
||||
rt, secret := meshRouter(t, c)
|
||||
h := rt.Handler()
|
||||
args := map[string]any{"model": "trellis2", "image": "data:image/png;base64," + testPicture(t), "seed": 3}
|
||||
|
||||
for _, streamed := range []bool{false, true} {
|
||||
res, events := mcp(t, h, secret, "tools/call", map[string]any{"name": "generate_mesh", "arguments": args,
|
||||
"_meta": map[string]any{"progressToken": "tok"}}, streamed)
|
||||
info := resultInfo(t, res)
|
||||
if info["status"] != "completed" || info["license"].(map[string]any)["name"] != "MIT" {
|
||||
t.Fatalf("streamed %v: %v", streamed, info)
|
||||
}
|
||||
id, _ := info["file_id"].(string)
|
||||
if !strings.HasSuffix(info["url"].(string), "/v1/files/"+id+"/content") {
|
||||
t.Errorf("url %v", info["url"])
|
||||
}
|
||||
// Fetched with the same key, it is the mesh.
|
||||
code, _, file := call(t, h, secret, http.MethodGet, "/v1/files/"+id+"/content", nil)
|
||||
if code != http.StatusOK || !strings.HasPrefix(file, "glTF") {
|
||||
t.Fatalf("the file is not the mesh: %d", code)
|
||||
}
|
||||
if streamed && !strings.Contains(strings.Join(events, "\n"), `"progressToken":"tok"`) {
|
||||
t.Errorf("no progress said while it ran: %v", events)
|
||||
}
|
||||
}
|
||||
// The job went through the client API as the caller: it is theirs.
|
||||
_, list, _ := call(t, h, secret, http.MethodGet, meshesPath, nil)
|
||||
if n := len(list["data"].([]any)); n != 2 {
|
||||
t.Errorf("%d jobs of the caller's, want the two made", n)
|
||||
}
|
||||
// Under no key, nothing is made.
|
||||
b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "tools/call",
|
||||
"params": map[string]any{"name": "generate_mesh", "arguments": args}})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(b))
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Errorf("a call with no key: %d", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPictureMadeThroughMCPComesWholeWhereNoFileCanBeKept(t *testing.T) {
|
||||
rt, queued, mu := imageRouter(t)
|
||||
h := rt.Handler()
|
||||
res, _ := mcp(t, h, "", "tools/call", map[string]any{"name": "generate_image",
|
||||
"arguments": map[string]any{"prompt": "a rock", "style": "concept", "size": 512, "seed": 9}}, false)
|
||||
info := resultInfo(t, res)
|
||||
if info["seed"] != 9.0 || info["model"] != "flux" || !strings.Contains(info["note"].(string), "keeps no files") {
|
||||
t.Errorf("info %v", info)
|
||||
}
|
||||
content := res["content"].([]any)
|
||||
if len(content) != 2 || content[1].(map[string]any)["type"] != "image" {
|
||||
t.Fatalf("the picture is not in the answer: %v", content)
|
||||
}
|
||||
if b, _ := base64.StdEncoding.DecodeString(content[1].(map[string]any)["data"].(string)); string(b) != "png-bytes" {
|
||||
t.Errorf("picture %q", b)
|
||||
}
|
||||
mu.Lock()
|
||||
wf, _ := json.Marshal(*queued)
|
||||
mu.Unlock()
|
||||
if !strings.Contains(string(wf), "a rock, single object centered") || !strings.Contains(string(wf), `"width":512`) {
|
||||
t.Errorf("the style was not added to the prompt: %s", wf)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPictureIsShrunkToItsLongSide(t *testing.T) {
|
||||
var b bytes.Buffer
|
||||
_ = png.Encode(&b, image.NewRGBA(image.Rect(0, 0, 2000, 1000)))
|
||||
out, err := shrink(b.Bytes(), 512)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
im, err := png.Decode(bytes.NewReader(out))
|
||||
if err != nil || im.Bounds().Dx() != 512 || im.Bounds().Dy() != 256 {
|
||||
t.Fatalf("%v %v", im.Bounds(), err)
|
||||
}
|
||||
b.Reset()
|
||||
_ = png.Encode(&b, image.NewRGBA(image.Rect(0, 0, 100, 50)))
|
||||
if out, _ := shrink(b.Bytes(), 512); !bytes.Equal(out, b.Bytes()) {
|
||||
t.Error("a small PNG was re-encoded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWorkflowRunsThroughMCPFromItsDeclaredInputs(t *testing.T) {
|
||||
wav := testWAV()
|
||||
c := newRunComfy(t, wav, "song_00001_.wav")
|
||||
rt, secret, _ := runRouter(t, c, songWorkflow, catalog.OutputAudio)
|
||||
h := rt.Handler()
|
||||
list, _ := mcp(t, h, secret, "tools/list", map[string]any{}, false)
|
||||
if !strings.Contains(strings.Join(toolNames(list), ","), "run_workflow") {
|
||||
t.Fatalf("no run_workflow for a model that declares its inputs: %v", toolNames(list))
|
||||
}
|
||||
res, _ := mcp(t, h, secret, "tools/call", map[string]any{"name": "run_workflow", "arguments": map[string]any{
|
||||
"model": "ace", "inputs": map[string]any{"tags": "synthwave", "cover": "data:image/png;base64," + testPicture(t)}}}, false)
|
||||
info := resultInfo(t, res)
|
||||
if info["status"] != "completed" || info["content_type"] != "audio/wav" {
|
||||
t.Fatalf("%v", info)
|
||||
}
|
||||
if _, _, file := call(t, h, secret, http.MethodGet, "/v1/files/"+info["file_id"].(string)+"/content", nil); file != string(wav) {
|
||||
t.Error("the file is not the song")
|
||||
}
|
||||
// A run refused says why, as the API did.
|
||||
res, _ = mcp(t, h, secret, "tools/call", map[string]any{"name": "run_workflow", "arguments": map[string]any{
|
||||
"model": "ace", "inputs": map[string]any{"cover": "data:image/png;base64," + testPicture(t)}}}, false)
|
||||
if res["isError"] != true || !strings.Contains(fmt.Sprint(res["content"]), "tags") {
|
||||
t.Errorf("a run missing a required input: %v", res)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPicturesAreLookedAtByAVisionModelShrunkAndToASchema(t *testing.T) {
|
||||
e := newFakeEngine(t, "vlm")
|
||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"vlm"}})
|
||||
vlm := catModel("vlm", nil)
|
||||
vlm.Capabilities.Vision = true
|
||||
cfg.Models = []catalog.Model{vlm}
|
||||
h := newTestRouter(t, cfg).Handler()
|
||||
|
||||
var big bytes.Buffer
|
||||
_ = png.Encode(&big, image.NewRGBA(image.Rect(0, 0, 3000, 1500)))
|
||||
res, _ := mcp(t, h, "", "tools/call", map[string]any{"name": "review_images", "arguments": map[string]any{
|
||||
"images": []string{base64.StdEncoding.EncodeToString(big.Bytes())},
|
||||
"question": "Does the render match the concept?",
|
||||
"json_schema": map[string]any{"type": "object", "properties": map[string]any{"match": map[string]any{"type": "boolean"}}},
|
||||
}}, false)
|
||||
if res["isError"] == true || !strings.Contains(fmt.Sprint(res["content"]), "hi") {
|
||||
t.Fatalf("%v", res)
|
||||
}
|
||||
var sent struct {
|
||||
ResponseFormat struct {
|
||||
Type string `json:"type"`
|
||||
} `json:"response_format"`
|
||||
Messages []struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
ImageURL struct {
|
||||
URL string `json:"url"`
|
||||
} `json:"image_url"`
|
||||
} `json:"content"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal(e.last.Load().([]byte), &sent); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sent.ResponseFormat.Type != "json_schema" || len(sent.Messages) != 1 || len(sent.Messages[0].Content) != 2 {
|
||||
t.Fatalf("sent %+v", sent)
|
||||
}
|
||||
b64, _ := strings.CutPrefix(sent.Messages[0].Content[1].ImageURL.URL, "data:image/png;base64,")
|
||||
pic, _ := base64.StdEncoding.DecodeString(b64)
|
||||
im, err := png.Decode(bytes.NewReader(pic))
|
||||
if err != nil || im.Bounds().Dx() != lookEdge || im.Bounds().Dy() != lookEdge/2 {
|
||||
t.Errorf("the picture sent is %v (%v), not shrunk to %d", im.Bounds(), err, lookEdge)
|
||||
}
|
||||
}
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
@@ -121,8 +120,7 @@ func (r *Router) handlePutNote(w http.ResponseWriter, req *http.Request) {
|
||||
Text string `json:"text"`
|
||||
Tags []string `json:"tags"`
|
||||
}
|
||||
// Pictures to look at come in the call, as base64.
|
||||
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, mcpBody))
|
||||
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
|
||||
if err != nil || json.Unmarshal(raw, &body) != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error",
|
||||
`body must be JSON like {"text": "...", "tags": ["..."]}`)
|
||||
@@ -275,66 +273,23 @@ func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
|
||||
resp := rpcResponse{JSONRPC: "2.0", ID: rpc.ID}
|
||||
switch rpc.Method {
|
||||
case "initialize":
|
||||
result := map[string]any{
|
||||
"protocolVersion": mcpVersionFor(rpc.Params),
|
||||
resp.Result = map[string]any{
|
||||
"protocolVersion": mcpProtocolVer,
|
||||
"capabilities": map[string]any{"tools": map[string]any{}},
|
||||
"serverInfo": map[string]any{"name": "llmgw", "version": "1"},
|
||||
"serverInfo": map[string]any{"name": "llmgw-memory", "version": "1"},
|
||||
}
|
||||
if r.mediaModels().any() {
|
||||
result["instructions"] = mcpMediaInstructions
|
||||
}
|
||||
resp.Result = result
|
||||
case "ping":
|
||||
resp.Result = map[string]any{}
|
||||
case "tools/list":
|
||||
resp.Result = map[string]any{"tools": append(slices.Clone(mcpTools), r.mediaTools()...)}
|
||||
resp.Result = map[string]any{"tools": mcpTools}
|
||||
case "tools/call":
|
||||
var call struct {
|
||||
Name string `json:"name"`
|
||||
Arguments json.RawMessage `json:"arguments"`
|
||||
Meta struct {
|
||||
ProgressToken json.RawMessage `json:"progressToken"`
|
||||
} `json:"_meta"`
|
||||
}
|
||||
_ = json.Unmarshal(rpc.Params, &call)
|
||||
if !isMediaTool(call.Name) {
|
||||
resp.Result = r.callMemoryTool(req, rpc.Params)
|
||||
break
|
||||
}
|
||||
if acceptsStream(req) {
|
||||
// Answered as a stream: how far along it is while it runs, which
|
||||
// a mesh's minutes want, then the answer.
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
sw := &streamWriter{w: w, token: call.Meta.ProgressToken}
|
||||
stop := make(chan struct{})
|
||||
go sw.keepAlive(stop, 15*time.Second)
|
||||
resp.Result = r.callMediaTool(mediaCall{req: req, progress: sw.progress, streamed: true}, call.Name, call.Arguments)
|
||||
close(stop)
|
||||
sw.event(resp)
|
||||
return
|
||||
}
|
||||
resp.Result = r.callMediaTool(mediaCall{req: req, progress: func(float64, string) {}}, call.Name, call.Arguments)
|
||||
resp.Result = r.callMemoryTool(req, rpc.Params)
|
||||
default:
|
||||
resp.Error = &rpcError{Code: -32601, Message: "method not found: " + rpc.Method}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// mcpVersionFor is the protocol version answered: the client's, of those
|
||||
// this server speaks, else its own.
|
||||
func mcpVersionFor(params json.RawMessage) string {
|
||||
var p struct {
|
||||
ProtocolVersion string `json:"protocolVersion"`
|
||||
}
|
||||
_ = json.Unmarshal(params, &p)
|
||||
if slices.Contains([]string{"2025-03-26", "2025-06-18"}, p.ProtocolVersion) {
|
||||
return p.ProtocolVersion
|
||||
}
|
||||
return mcpProtocolVer
|
||||
}
|
||||
|
||||
func toolText(text string, isError bool) map[string]any {
|
||||
return map[string]any{
|
||||
"content": []map[string]any{{"type": "text", "text": text}},
|
||||
|
||||
@@ -51,10 +51,6 @@ type Router struct {
|
||||
cfg config.Router
|
||||
log *slog.Logger
|
||||
|
||||
// selfHandler is the client API as Handler built it, key check and all:
|
||||
// what the MCP tools make their requests through, as the caller.
|
||||
selfHandler atomic.Pointer[http.Handler]
|
||||
|
||||
reg *endpoint.Registry
|
||||
poller *scrape.Poller
|
||||
px *proxy.Proxy
|
||||
|
||||
Reference in new issue
Block a user