feat: add brain_query and brain_write MCP tools
Adds the brain skill that proxies HTTP calls to the ingestion server, exposing brain_query (/query) and brain_write (/write) as MCP tools. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
90
internal/skills/brain/handlers.go
Normal file
90
internal/skills/brain/handlers.go
Normal file
@@ -0,0 +1,90 @@
|
|||||||
|
// internal/skills/brain/handlers.go
|
||||||
|
package brain
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Handle dispatches brain_query and brain_write tool calls.
|
||||||
|
func (s *Skill) Handle(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
|
||||||
|
switch tool {
|
||||||
|
case "brain_query":
|
||||||
|
return s.query(ctx, args)
|
||||||
|
case "brain_write":
|
||||||
|
return s.write(ctx, args)
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unknown brain tool: %s", tool)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type queryArgs struct {
|
||||||
|
Query string `json:"query"`
|
||||||
|
Limit int `json:"limit,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Skill) query(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
|
||||||
|
var a queryArgs
|
||||||
|
if err := json.Unmarshal(args, &a); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse args: %w", err)
|
||||||
|
}
|
||||||
|
if a.Query == "" {
|
||||||
|
return nil, fmt.Errorf("query is required")
|
||||||
|
}
|
||||||
|
if a.Limit == 0 {
|
||||||
|
a.Limit = 5
|
||||||
|
}
|
||||||
|
return s.post(ctx, "/query", a)
|
||||||
|
}
|
||||||
|
|
||||||
|
type writeArgs struct {
|
||||||
|
Content string `json:"content"`
|
||||||
|
Type string `json:"type,omitempty"`
|
||||||
|
Domain string `json:"domain,omitempty"`
|
||||||
|
Filename string `json:"filename,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Skill) write(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
|
||||||
|
var a writeArgs
|
||||||
|
if err := json.Unmarshal(args, &a); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse args: %w", err)
|
||||||
|
}
|
||||||
|
if a.Content == "" {
|
||||||
|
return nil, fmt.Errorf("content is required")
|
||||||
|
}
|
||||||
|
return s.post(ctx, "/write", map[string]string{
|
||||||
|
"content": a.Content,
|
||||||
|
"filename": a.Filename,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Skill) post(ctx context.Context, path string, body any) (json.RawMessage, error) {
|
||||||
|
b, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("marshal request: %w", err)
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.cfg.IngestBaseURL+path, bytes.NewReader(b))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("call ingestion server: %w", err)
|
||||||
|
}
|
||||||
|
defer func() { _ = resp.Body.Close() }()
|
||||||
|
|
||||||
|
out, err := io.ReadAll(resp.Body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("read response: %w", err)
|
||||||
|
}
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return nil, fmt.Errorf("ingestion server returned %d: %s", resp.StatusCode, out)
|
||||||
|
}
|
||||||
|
return json.RawMessage(out), nil
|
||||||
|
}
|
||||||
61
internal/skills/brain/handlers_test.go
Normal file
61
internal/skills/brain/handlers_test.go
Normal file
@@ -0,0 +1,61 @@
|
|||||||
|
// internal/skills/brain/handlers_test.go
|
||||||
|
package brain_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/mathiasbq/supervisor/internal/skills/brain"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestHandle_BrainQuery_CallsIngestServer(t *testing.T) {
|
||||||
|
called := false
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
assert.Equal(t, "/query", r.URL.Path)
|
||||||
|
called = true
|
||||||
|
json.NewEncoder(w).Encode(map[string]any{
|
||||||
|
"results": []map[string]any{
|
||||||
|
{"path": "wiki/concepts/tdd.md", "title": "TDD", "excerpt": "Test-driven development.", "score": 3},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
s := brain.New(brain.Config{IngestBaseURL: srv.URL})
|
||||||
|
args, _ := json.Marshal(map[string]string{"query": "test driven development"})
|
||||||
|
out, err := s.Handle(context.Background(), "brain_query", args)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.True(t, called)
|
||||||
|
|
||||||
|
var result map[string]any
|
||||||
|
require.NoError(t, json.Unmarshal(out, &result))
|
||||||
|
results := result["results"].([]any)
|
||||||
|
assert.Len(t, results, 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandle_BrainWrite_CallsIngestServer(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
assert.Equal(t, "/write", r.URL.Path)
|
||||||
|
json.NewEncoder(w).Encode(map[string]string{"path": "raw/test.md"})
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
s := brain.New(brain.Config{IngestBaseURL: srv.URL})
|
||||||
|
args, _ := json.Marshal(map[string]string{"content": "# Test\n\nSome learning.", "type": "concept"})
|
||||||
|
out, err := s.Handle(context.Background(), "brain_write", args)
|
||||||
|
require.NoError(t, err)
|
||||||
|
var result map[string]string
|
||||||
|
require.NoError(t, json.Unmarshal(out, &result))
|
||||||
|
assert.Equal(t, "raw/test.md", result["path"])
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandle_UnknownTool_ReturnsError(t *testing.T) {
|
||||||
|
s := brain.New(brain.Config{IngestBaseURL: "http://localhost:3300"})
|
||||||
|
_, err := s.Handle(context.Background(), "brain_unknown", nil)
|
||||||
|
assert.Error(t, err)
|
||||||
|
}
|
||||||
55
internal/skills/brain/skill.go
Normal file
55
internal/skills/brain/skill.go
Normal file
@@ -0,0 +1,55 @@
|
|||||||
|
// internal/skills/brain/skill.go
|
||||||
|
package brain
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
|
||||||
|
"github.com/mathiasbq/supervisor/internal/registry"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Config holds brain skill configuration.
|
||||||
|
type Config struct {
|
||||||
|
IngestBaseURL string // base URL of the ingestion HTTP server, e.g. http://localhost:3300
|
||||||
|
}
|
||||||
|
|
||||||
|
// Skill implements registry.Skill for brain_query and brain_write.
|
||||||
|
type Skill struct {
|
||||||
|
cfg Config
|
||||||
|
}
|
||||||
|
|
||||||
|
// New constructs a brain Skill.
|
||||||
|
func New(cfg Config) *Skill { return &Skill{cfg: cfg} }
|
||||||
|
|
||||||
|
// Name returns the skill name used for routing.
|
||||||
|
func (s *Skill) Name() string { return "brain" }
|
||||||
|
|
||||||
|
// Tools returns the MCP tool definitions for brain_query and brain_write.
|
||||||
|
func (s *Skill) Tools() []registry.ToolDef {
|
||||||
|
schema := func(required []string, props map[string]any) json.RawMessage {
|
||||||
|
b, _ := json.Marshal(map[string]any{"type": "object", "required": required, "properties": props})
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
str := map[string]any{"type": "string"}
|
||||||
|
num := map[string]any{"type": "integer"}
|
||||||
|
|
||||||
|
return []registry.ToolDef{
|
||||||
|
{
|
||||||
|
Name: "brain_query",
|
||||||
|
Description: "Search the hyperguild brain wiki for relevant knowledge. Call this before starting any significant task.",
|
||||||
|
InputSchema: schema([]string{"query"}, map[string]any{
|
||||||
|
"query": str,
|
||||||
|
"limit": num,
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "brain_write",
|
||||||
|
Description: "Write a raw knowledge note to the brain for later ingestion into the wiki.",
|
||||||
|
InputSchema: schema([]string{"content"}, map[string]any{
|
||||||
|
"content": str,
|
||||||
|
"type": str,
|
||||||
|
"domain": str,
|
||||||
|
"filename": str,
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user