mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-29 17:44:30 -04:00
feat(failover): add MCP tools to list chains and pin targets
Assisted-by: Claude:claude-opus-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
ea55e2ffc3
commit
191d2d6301
19 files changed
+536
-20
No files matched your search
@@ -130,4 +130,15 @@ type LocalAIClient interface {
|
||||
// ClearRouterCorpus wipes a knn router's corpus — file and live
|
||||
// index.
|
||||
ClearRouterCorpus(ctx context.Context, routerModel string) (*RouterCorpusClearResult, error)
|
||||
|
||||
// ---- Failover chains ----
|
||||
// ListFailoverChains reports every configured failover chain, its
|
||||
// currently active target, and the health of each target.
|
||||
ListFailoverChains(ctx context.Context) ([]FailoverChainInfo, error)
|
||||
// PinFailoverTarget forces chain to serve every request from target,
|
||||
// regardless of health, until unpinned.
|
||||
PinFailoverTarget(ctx context.Context, chain, target string) error
|
||||
// UnpinFailoverTarget removes chain's pin so health decides the
|
||||
// active target again.
|
||||
UnpinFailoverTarget(ctx context.Context, chain string) error
|
||||
}
|
||||
@@ -46,27 +46,30 @@ var toolToHTTPRoute = map[string]string{
|
||||
ToolGetRouterCorpusStats: "GET /api/router/:name/corpus/stats",
|
||||
ToolListAliases: "GET /api/aliases",
|
||||
ToolListVoiceProfiles: "GET /api/voice-profiles",
|
||||
ToolListFailoverChains: "GET /api/failover",
|
||||
|
||||
// Mutating tools.
|
||||
ToolInstallModel: "POST /models/apply",
|
||||
ToolImportModelURI: "POST /models/import-uri",
|
||||
ToolDeleteModel: "POST /models/delete/:name",
|
||||
ToolEditModelConfig: "PATCH /api/models/config-json/:name",
|
||||
ToolReloadModels: "POST /models/reload",
|
||||
ToolLoadModel: "POST /backend/load",
|
||||
ToolInstallBackend: "POST /backends/apply",
|
||||
ToolUpgradeBackend: "POST /backends/upgrade/:name",
|
||||
ToolToggleModelState: "PUT /models/toggle-state/:name/:action",
|
||||
ToolToggleModelPinned: "PUT /models/toggle-pinned/:name/:action",
|
||||
ToolSetBranding: "POST /api/settings (instance_name, instance_tagline)",
|
||||
ToolSetAlias: "PATCH /api/models/config-json/:name (swap) or POST /models/import (create)",
|
||||
ToolSeedRouterCorpus: "POST /api/router/:name/corpus",
|
||||
ToolClearRouterCorpus: "DELETE /api/router/:name/corpus",
|
||||
ToolCreateVoiceProfile: "POST /api/voice-profiles",
|
||||
ToolDeleteVoiceProfile: "DELETE /api/voice-profiles/:id",
|
||||
ToolSetNodeVRAMBudget: "PUT /api/nodes/:id/vram-budget",
|
||||
ToolSetScheduling: "POST /api/nodes/scheduling",
|
||||
ToolDeleteScheduling: "DELETE /api/nodes/scheduling/:model",
|
||||
ToolInstallModel: "POST /models/apply",
|
||||
ToolImportModelURI: "POST /models/import-uri",
|
||||
ToolDeleteModel: "POST /models/delete/:name",
|
||||
ToolEditModelConfig: "PATCH /api/models/config-json/:name",
|
||||
ToolReloadModels: "POST /models/reload",
|
||||
ToolLoadModel: "POST /backend/load",
|
||||
ToolInstallBackend: "POST /backends/apply",
|
||||
ToolUpgradeBackend: "POST /backends/upgrade/:name",
|
||||
ToolToggleModelState: "PUT /models/toggle-state/:name/:action",
|
||||
ToolToggleModelPinned: "PUT /models/toggle-pinned/:name/:action",
|
||||
ToolSetBranding: "POST /api/settings (instance_name, instance_tagline)",
|
||||
ToolSetAlias: "PATCH /api/models/config-json/:name (swap) or POST /models/import (create)",
|
||||
ToolSeedRouterCorpus: "POST /api/router/:name/corpus",
|
||||
ToolClearRouterCorpus: "DELETE /api/router/:name/corpus",
|
||||
ToolCreateVoiceProfile: "POST /api/voice-profiles",
|
||||
ToolDeleteVoiceProfile: "DELETE /api/voice-profiles/:id",
|
||||
ToolSetNodeVRAMBudget: "PUT /api/nodes/:id/vram-budget",
|
||||
ToolSetScheduling: "POST /api/nodes/scheduling",
|
||||
ToolDeleteScheduling: "DELETE /api/nodes/scheduling/:model",
|
||||
ToolPinFailoverTarget: "POST /api/failover/:chain/pin",
|
||||
ToolUnpinFailoverTarget: "DELETE /api/failover/:chain/pin",
|
||||
}
|
||||
|
||||
// allKnownTools is the union of expectedFullCatalog (defined in
|
||||
|
||||
@@ -413,3 +413,26 @@ type VRAMEstimateRequest struct {
|
||||
GPULayers int `json:"gpu_layers,omitempty" jsonschema:"Number of layers to offload to GPU. -1 for all."`
|
||||
KVQuantBits int `json:"kv_quant_bits,omitempty" jsonschema:"KV cache quantization bits (e.g. 4, 8, 16)."`
|
||||
}
|
||||
|
||||
// FailoverTargetInfo is the LLM-facing view of one failover chain target's
|
||||
// health. It mirrors failover.TargetStatus but drops ConsecutiveOK and
|
||||
// LastProbe — internal probing detail the LLM doesn't need to decide
|
||||
// whether to pin or unpin a target.
|
||||
type FailoverTargetInfo struct {
|
||||
Model string `json:"model"`
|
||||
Kind string `json:"kind"`
|
||||
Warm bool `json:"warm"`
|
||||
State string `json:"state"`
|
||||
LastError string `json:"last_error,omitempty"`
|
||||
}
|
||||
|
||||
// FailoverChainInfo is the LLM-facing view of one failover chain: its
|
||||
// current state, the target serving it now, an optional pin, and every
|
||||
// target's health.
|
||||
type FailoverChainInfo struct {
|
||||
Name string `json:"name"`
|
||||
State string `json:"state"`
|
||||
Active string `json:"active"`
|
||||
Pinned string `json:"pinned,omitempty"`
|
||||
Targets []FailoverTargetInfo `json:"targets"`
|
||||
}
|
||||
@@ -59,6 +59,9 @@ type fakeClient struct {
|
||||
getPIIEvents func(PIIEventsQuery) ([]PIIEvent, error)
|
||||
getMiddlewareStatus func() (*MiddlewareStatus, error)
|
||||
getRouterDecisions func(RouterDecisionsQuery) ([]RouterDecision, error)
|
||||
listFailoverChains func() ([]FailoverChainInfo, error)
|
||||
pinFailoverTarget func(string, string) error
|
||||
unpinFailoverTarget func(string) error
|
||||
}
|
||||
|
||||
type fakeCall struct {
|
||||
@@ -393,3 +396,27 @@ func (f *fakeClient) ClearRouterCorpus(_ context.Context, routerModel string) (*
|
||||
f.record("ClearRouterCorpus", routerModel)
|
||||
return &RouterCorpusClearResult{Router: routerModel}, nil
|
||||
}
|
||||
|
||||
func (f *fakeClient) ListFailoverChains(_ context.Context) ([]FailoverChainInfo, error) {
|
||||
f.record("ListFailoverChains", nil)
|
||||
if f.listFailoverChains != nil {
|
||||
return f.listFailoverChains()
|
||||
}
|
||||
return []FailoverChainInfo{}, nil
|
||||
}
|
||||
|
||||
func (f *fakeClient) PinFailoverTarget(_ context.Context, chain, target string) error {
|
||||
f.record("PinFailoverTarget", []any{chain, target})
|
||||
if f.pinFailoverTarget != nil {
|
||||
return f.pinFailoverTarget(chain, target)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeClient) UnpinFailoverTarget(_ context.Context, chain string) error {
|
||||
f.record("UnpinFailoverTarget", chain)
|
||||
if f.unpinFailoverTarget != nil {
|
||||
return f.unpinFailoverTarget(chain)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -829,3 +829,23 @@ func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*lo
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// ---- Failover chains ----
|
||||
|
||||
func (c *Client) ListFailoverChains(ctx context.Context) ([]localaitools.FailoverChainInfo, error) {
|
||||
var out struct {
|
||||
Chains []localaitools.FailoverChainInfo `json:"chains"`
|
||||
}
|
||||
if err := c.do(ctx, http.MethodGet, routeFailover, nil, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out.Chains, nil
|
||||
}
|
||||
|
||||
func (c *Client) PinFailoverTarget(ctx context.Context, chain, target string) error {
|
||||
return c.do(ctx, http.MethodPost, routeFailover+"/"+url.PathEscape(chain)+"/pin", map[string]string{"target": target}, nil)
|
||||
}
|
||||
|
||||
func (c *Client) UnpinFailoverTarget(ctx context.Context, chain string) error {
|
||||
return c.do(ctx, http.MethodDelete, routeFailover+"/"+url.PathEscape(chain)+"/pin", nil, nil)
|
||||
}
|
||||
@@ -365,6 +365,68 @@ var _ = Describe("Model aliases", func() {
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Failover chains", func() {
|
||||
Describe("ListFailoverChains", func() {
|
||||
It("issues GET /api/failover and unwraps the chains array", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
Expect(r.Method).To(Equal(http.MethodGet))
|
||||
Expect(r.URL.Path).To(Equal("/api/failover"))
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"chains": []map[string]any{
|
||||
{
|
||||
"name": "chain",
|
||||
"state": "primary",
|
||||
"active": "a",
|
||||
"pinned": nil,
|
||||
"targets": []map[string]any{
|
||||
{"model": "a", "kind": "local", "warm": false, "state": "healthy"},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
}))
|
||||
DeferCleanup(srv.Close)
|
||||
|
||||
out, err := New(srv.URL, "").ListFailoverChains(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(HaveLen(1))
|
||||
Expect(out[0].Name).To(Equal("chain"))
|
||||
Expect(out[0].Active).To(Equal("a"))
|
||||
Expect(out[0].Pinned).To(BeEmpty())
|
||||
Expect(out[0].Targets).To(ConsistOf(localaitools.FailoverTargetInfo{Model: "a", Kind: "local", State: "healthy"}))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("PinFailoverTarget", func() {
|
||||
It("issues POST /api/failover/:chain/pin with the target in the body", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
Expect(r.Method).To(Equal(http.MethodPost))
|
||||
Expect(r.URL.Path).To(Equal("/api/failover/chain/pin"))
|
||||
var body map[string]string
|
||||
Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed())
|
||||
Expect(body).To(HaveKeyWithValue("target", "b"))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
DeferCleanup(srv.Close)
|
||||
|
||||
Expect(New(srv.URL, "").PinFailoverTarget(context.Background(), "chain", "b")).To(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("UnpinFailoverTarget", func() {
|
||||
It("issues DELETE /api/failover/:chain/pin", func() {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
Expect(r.Method).To(Equal(http.MethodDelete))
|
||||
Expect(r.URL.Path).To(Equal("/api/failover/chain/pin"))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
DeferCleanup(srv.Close)
|
||||
|
||||
Expect(New(srv.URL, "").UnpinFailoverTarget(context.Background(), "chain")).To(Succeed())
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("ErrHTTPNotFound", func() {
|
||||
Context("on a clean 404 status", func() {
|
||||
var (
|
||||
|
||||
@@ -34,6 +34,7 @@ const (
|
||||
routeMiddleware = "/api/middleware/status"
|
||||
routeRouterDecisions = "/api/router/decisions"
|
||||
routeVoiceProfiles = "/api/voice-profiles"
|
||||
routeFailover = "/api/failover"
|
||||
)
|
||||
|
||||
func routeJobStatus(jobID string) string {
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
"github.com/mudler/LocalAI/core/gallery/importers"
|
||||
"github.com/mudler/LocalAI/core/http/auth"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/failover"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/core/services/modeladmin"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
@@ -83,6 +84,12 @@ type Client struct {
|
||||
RouterEmbedderFingerprint func(modelName string) (string, error)
|
||||
RouterVectorStore func(storeName string) backend.VectorStore
|
||||
|
||||
// Failover backs list_failover_chains / pin_failover_target /
|
||||
// unpin_failover_target. nil makes the tools report "failover is not
|
||||
// running" — the same as a deployment with no failover chains
|
||||
// configured.
|
||||
Failover *failover.Manager
|
||||
|
||||
modelAdmin *modeladmin.ConfigService
|
||||
}
|
||||
|
||||
@@ -1103,3 +1110,39 @@ func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*lo
|
||||
}
|
||||
return &localaitools.RouterCorpusClearResult{Router: cfg.Name, Cleared: cleared}, nil
|
||||
}
|
||||
|
||||
// ---- Failover chains ----
|
||||
|
||||
func (c *Client) ListFailoverChains(_ context.Context) ([]localaitools.FailoverChainInfo, error) {
|
||||
out := []localaitools.FailoverChainInfo{}
|
||||
if c.Failover == nil {
|
||||
return out, nil
|
||||
}
|
||||
for _, ch := range c.Failover.Status() {
|
||||
info := localaitools.FailoverChainInfo{Name: ch.Name, State: string(ch.State), Active: ch.Active}
|
||||
if ch.Pinned != nil {
|
||||
info.Pinned = *ch.Pinned
|
||||
}
|
||||
for _, t := range ch.Targets {
|
||||
info.Targets = append(info.Targets, localaitools.FailoverTargetInfo{
|
||||
Model: t.Model, Kind: string(t.Kind), Warm: t.Warm, State: string(t.State), LastError: t.LastError,
|
||||
})
|
||||
}
|
||||
out = append(out, info)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *Client) PinFailoverTarget(_ context.Context, chain, target string) error {
|
||||
if c.Failover == nil {
|
||||
return errors.New("failover is not running")
|
||||
}
|
||||
return c.Failover.Pin(chain, target)
|
||||
}
|
||||
|
||||
func (c *Client) UnpinFailoverTarget(_ context.Context, chain string) error {
|
||||
if c.Failover == nil {
|
||||
return errors.New("failover is not running")
|
||||
}
|
||||
return c.Failover.Unpin(chain)
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/core/services/failover"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
localaitools "github.com/mudler/LocalAI/pkg/mcp/localaitools"
|
||||
@@ -129,6 +130,95 @@ var _ = Describe("inproc.Client model aliases", func() {
|
||||
})
|
||||
})
|
||||
|
||||
// fakeFailoverSource is a minimal failover.ConfigSource over an in-memory
|
||||
// map, so these specs don't need a real ModelConfigLoader + on-disk YAML.
|
||||
type fakeFailoverSource struct {
|
||||
cfgs map[string]config.ModelConfig
|
||||
}
|
||||
|
||||
func (s *fakeFailoverSource) GetModelConfig(name string) (config.ModelConfig, bool) {
|
||||
c, ok := s.cfgs[name]
|
||||
return c, ok
|
||||
}
|
||||
|
||||
func (s *fakeFailoverSource) GetAllModelsConfigs() []config.ModelConfig {
|
||||
out := make([]config.ModelConfig, 0, len(s.cfgs))
|
||||
for _, c := range s.cfgs {
|
||||
out = append(out, c)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
var _ = Describe("inproc.Client failover chains", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
c *Client
|
||||
fm *failover.Manager
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx = context.Background()
|
||||
src := &fakeFailoverSource{cfgs: map[string]config.ModelConfig{
|
||||
"a": {Name: "a", Backend: "llama-cpp"},
|
||||
"b": {Name: "b", Backend: "llama-cpp"},
|
||||
"chain": {Name: "chain", Failover: &config.FailoverConfig{
|
||||
Targets: []config.FailoverTarget{{Model: "a"}, {Model: "b"}},
|
||||
}},
|
||||
}}
|
||||
fm = failover.New(src)
|
||||
c = &Client{Failover: fm}
|
||||
})
|
||||
|
||||
It("ListFailoverChains reports the chain, its active target, and target health", func() {
|
||||
out, err := c.ListFailoverChains(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(HaveLen(1))
|
||||
Expect(out[0].Name).To(Equal("chain"))
|
||||
Expect(out[0].Active).To(Equal("a"))
|
||||
Expect(out[0].Pinned).To(BeEmpty())
|
||||
Expect(out[0].Targets).To(HaveLen(2))
|
||||
Expect(out[0].Targets[0].Model).To(Equal("a"))
|
||||
Expect(out[0].Targets[0].Kind).To(Equal("local"))
|
||||
Expect(out[0].Targets[0].State).To(Equal("healthy"))
|
||||
})
|
||||
|
||||
It("returns an empty slice, not an error, when no failover manager is wired", func() {
|
||||
c = &Client{}
|
||||
out, err := c.ListFailoverChains(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("PinFailoverTarget pins the chain and ListFailoverChains reflects it", func() {
|
||||
Expect(c.PinFailoverTarget(ctx, "chain", "b")).To(Succeed())
|
||||
|
||||
out, err := c.ListFailoverChains(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out[0].Pinned).To(Equal("b"))
|
||||
})
|
||||
|
||||
It("PinFailoverTarget errors when the failover manager is unavailable", func() {
|
||||
c = &Client{}
|
||||
err := c.PinFailoverTarget(ctx, "chain", "b")
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("UnpinFailoverTarget clears a pin", func() {
|
||||
Expect(c.PinFailoverTarget(ctx, "chain", "b")).To(Succeed())
|
||||
Expect(c.UnpinFailoverTarget(ctx, "chain")).To(Succeed())
|
||||
|
||||
out, err := c.ListFailoverChains(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out[0].Pinned).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("UnpinFailoverTarget errors when the failover manager is unavailable", func() {
|
||||
c = &Client{}
|
||||
err := c.UnpinFailoverTarget(ctx, "chain")
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("inproc.Client model scheduling", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
These rules are non-negotiable. The user trusts you to operate their server without unintended changes.
|
||||
|
||||
1. **Confirm before mutating.** Before calling any of these tools — `install_model`, `import_model_uri`, `delete_model`, `install_backend`, `upgrade_backend`, `edit_model_config`, `reload_models`, `load_model`, `toggle_model_state`, `toggle_model_pinned`, `set_branding`, `set_alias`, `seed_router_corpus`, `clear_router_corpus`, `create_voice_profile`, `delete_voice_profile`, `set_node_vram_budget`, `set_scheduling`, `delete_scheduling` — first state in plain language what you are about to do (which tool, which target, which arguments) and wait for the user's explicit confirmation in the next turn. "Yes", "do it", "go ahead", "proceed" all count as confirmation. Anything else does not.
|
||||
1. **Confirm before mutating.** Before calling any of these tools — `install_model`, `import_model_uri`, `delete_model`, `install_backend`, `upgrade_backend`, `edit_model_config`, `reload_models`, `load_model`, `toggle_model_state`, `toggle_model_pinned`, `set_branding`, `set_alias`, `seed_router_corpus`, `clear_router_corpus`, `create_voice_profile`, `delete_voice_profile`, `set_node_vram_budget`, `set_scheduling`, `delete_scheduling`, `pin_failover_target`, `unpin_failover_target` — first state in plain language what you are about to do (which tool, which target, which arguments) and wait for the user's explicit confirmation in the next turn. "Yes", "do it", "go ahead", "proceed" all count as confirmation. Anything else does not.
|
||||
|
||||
2. **Disambiguate before mutating.** If the user's request is ambiguous (several gallery candidates match, the model name has multiple installed versions, the backend has variants), present the candidates as a numbered list and ask the user to pick before calling any mutating tool.
|
||||
|
||||
|
||||
@@ -24,6 +24,7 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the
|
||||
- `get_router_decisions` — Inspect recent router decisions and classifier signals.
|
||||
- `get_router_corpus_stats` — Inspect a KNN router corpus by count and label only; exemplar texts are never returned.
|
||||
- `list_aliases` — List configured model aliases and their targets.
|
||||
- `list_failover_chains` — List failover chains, their active target and target health.
|
||||
|
||||
## Mutating (require user confirmation per safety rule 1)
|
||||
|
||||
@@ -46,3 +47,5 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the
|
||||
- `set_node_vram_budget` — Set or clear a federated node's VRAM budget override.
|
||||
- `set_scheduling` — Create or update a distributed per-model scheduling config.
|
||||
- `delete_scheduling` — Remove a distributed per-model scheduling config.
|
||||
- `pin_failover_target` — Force a failover chain to one target.
|
||||
- `unpin_failover_target` — Remove a failover pin.
|
||||
@@ -54,6 +54,7 @@ func NewServer(client LocalAIClient, opts Options) *mcp.Server {
|
||||
registerUsageTools(srv, client, opts)
|
||||
registerPIITools(srv, client, opts)
|
||||
registerMiddlewareTools(srv, client, opts)
|
||||
registerFailoverTools(srv, client, opts)
|
||||
|
||||
return srv
|
||||
}
|
||||
@@ -92,6 +92,7 @@ var expectedReadOnlyCatalog = sortedStrings(
|
||||
ToolListVoiceProfiles,
|
||||
ToolSystemInfo,
|
||||
ToolVRAMEstimate,
|
||||
ToolListFailoverChains,
|
||||
)
|
||||
|
||||
// expectedFullCatalog derives from the read-only catalog plus the canonical
|
||||
@@ -155,6 +156,9 @@ var _ = Describe("Tool dispatch", func() {
|
||||
{ToolListAliases, struct{}{}, "ListAliases"},
|
||||
{ToolCreateVoiceProfile, CreateVoiceProfileRequest{Name: "Narrator", Transcript: "Reference words", AudioBase64: "UklGRg==", ConsentConfirmed: true}, "CreateVoiceProfile"},
|
||||
{ToolDeleteVoiceProfile, DeleteVoiceProfileRequest{ID: "00000000-0000-0000-0000-000000000001"}, "DeleteVoiceProfile"},
|
||||
{ToolListFailoverChains, map[string]any{}, "ListFailoverChains"},
|
||||
{ToolPinFailoverTarget, map[string]any{"chain": "c", "target": "b"}, "PinFailoverTarget"},
|
||||
{ToolUnpinFailoverTarget, map[string]any{"chain": "c"}, "UnpinFailoverTarget"},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
|
||||
@@ -49,10 +49,19 @@ const (
|
||||
ToolSetNodeVRAMBudget = "set_node_vram_budget"
|
||||
ToolSetScheduling = "set_scheduling"
|
||||
ToolDeleteScheduling = "delete_scheduling"
|
||||
// ToolPinFailoverTarget and ToolUnpinFailoverTarget live here (rather
|
||||
// than grouped with ToolListFailoverChains below) so mutatingToolNames
|
||||
// stays a contiguous scan of this block.
|
||||
ToolPinFailoverTarget = "pin_failover_target"
|
||||
ToolUnpinFailoverTarget = "unpin_failover_target"
|
||||
|
||||
// ToolListAliases is read-only but lives here so the alias tools stay
|
||||
// grouped; the catalog tests assert its read-only placement.
|
||||
ToolListAliases = "list_aliases"
|
||||
|
||||
// ToolListFailoverChains is read-only but lives here so the failover
|
||||
// tools stay grouped; the catalog tests assert its read-only placement.
|
||||
ToolListFailoverChains = "list_failover_chains"
|
||||
)
|
||||
|
||||
// DefaultServerName is the MCP Implementation.Name surfaced when
|
||||
@@ -83,4 +92,6 @@ var mutatingToolNames = []string{
|
||||
ToolSetNodeVRAMBudget,
|
||||
ToolSetScheduling,
|
||||
ToolDeleteScheduling,
|
||||
ToolPinFailoverTarget,
|
||||
ToolUnpinFailoverTarget,
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package localaitools
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
// registerFailoverTools wires the conversational failover-chain tools.
|
||||
// list_failover_chains reports the health of every chain, pin_failover_target
|
||||
// forces a chain to one target, and unpin_failover_target hands control back
|
||||
// to health-based selection.
|
||||
func registerFailoverTools(s *mcp.Server, client LocalAIClient, opts Options) {
|
||||
mcp.AddTool(s, &mcp.Tool{
|
||||
Name: ToolListFailoverChains,
|
||||
Description: "List model failover chains, the target serving each one now, and the health of every target.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) {
|
||||
chains, err := client.ListFailoverChains(ctx)
|
||||
if err != nil {
|
||||
return errorResult(err), nil, nil
|
||||
}
|
||||
return jsonResult(chains), nil, nil
|
||||
})
|
||||
|
||||
if opts.DisableMutating {
|
||||
return
|
||||
}
|
||||
|
||||
mcp.AddTool(s, &mcp.Tool{
|
||||
Name: ToolPinFailoverTarget,
|
||||
Description: "Force a failover chain to serve every request from one target, regardless of health, until it is unpinned. Requires user confirmation per safety rule 1.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, args struct {
|
||||
Chain string `json:"chain" jsonschema:"failover chain name"`
|
||||
Target string `json:"target" jsonschema:"target model to pin"`
|
||||
}) (*mcp.CallToolResult, any, error) {
|
||||
if err := client.PinFailoverTarget(ctx, args.Chain, args.Target); err != nil {
|
||||
return errorResult(err), nil, nil
|
||||
}
|
||||
return jsonResult(map[string]string{"chain": args.Chain, "pinned": args.Target}), nil, nil
|
||||
})
|
||||
|
||||
mcp.AddTool(s, &mcp.Tool{
|
||||
Name: ToolUnpinFailoverTarget,
|
||||
Description: "Remove the pin from a failover chain so health decides the target again. Requires user confirmation per safety rule 1.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, args struct {
|
||||
Chain string `json:"chain" jsonschema:"failover chain name"`
|
||||
}) (*mcp.CallToolResult, any, error) {
|
||||
if err := client.UnpinFailoverTarget(ctx, args.Chain); err != nil {
|
||||
return errorResult(err), nil, nil
|
||||
}
|
||||
return jsonResult(map[string]string{"chain": args.Chain, "pinned": ""}), nil, nil
|
||||
})
|
||||
}
|
||||
Reference in new issue
Block a user