feat(ai): invoke a project skill with / in the analyst chat

runtime/ai/analyst_agent.go

@@ -182,6 +182,19 @@ func (t *AnalystAgent) Handler(ctx context.Context, args *AnalystAgentArgs) (*An
 		}
 	}
 
+	// Pre-invoke the load_skill tool for each analyst skill the user referenced in the prompt, on every invocation.
+	// A skill already loaded in this conversation is skipped: the model has it, and loading it again would repeat its whole body in the context.
+	loaded := loadedSkills(s)
+	for _, sk := range referencedSkills(args.Prompt, skillsForAgent(skills, parser.SkillAgentAnalyst)) {
+		if loaded[sk.Name] {
+			continue
+		}
+		_, err := s.CallTool(ctx, RoleAssistant, LoadSkillName, nil, &LoadSkillArgs{Name: sk.Name})
+		if err != nil && errors.Is(err, ctx.Err()) { // Don't exit on non-context errors
+			return nil, err
+		}
+	}
+
 	// Determine tools that can be used
 	tools := []string{}
 	if args.Explore == "" {

runtime/ai/skill_references.go

@@ -0,0 +1,50 @@
+package ai
+
+import (
+	"regexp"
+	"slices"
+)
+
+var (
+	// chatReferenceRegexp matches the references that the chat UI writes into prompts, e.g. <chat-reference>type="skill" skill="monthly-close"</chat-reference>.
+	chatReferenceRegexp = regexp.MustCompile(`(?s)<chat-reference\b(.*?)</chat-reference>`)
+	// chatReferenceAttrRegexp matches a key="value" pair in a chat reference.
+	chatReferenceAttrRegexp = regexp.MustCompile(`(\w+)="([^"]*)"`)
+)
+
+// referencedSkills returns the skills referenced in a prompt with a chat reference of type "skill", once each and in order of first reference.
+// References to skills that are not in the given list are ignored.
+func referencedSkills(prompt string, skills []*Skill) []*Skill {
+	var res []*Skill
+	for _, ref := range chatReferenceRegexp.FindAllStringSubmatch(prompt, -1) {
+		attrs := map[string]string{}
+		for _, attr := range chatReferenceAttrRegexp.FindAllStringSubmatch(ref[1], -1) {
+			attrs[attr[1]] = attr[2]
+		}
+		if attrs["type"] != "skill" {
+			continue
+		}
+
+		idx := slices.IndexFunc(skills, func(sk *Skill) bool { return sk.Name == attrs["skill"] })
+		if idx == -1 || slices.Contains(res, skills[idx]) {
+			continue
+		}
+		res = append(res, skills[idx])
+	}
+	return res
+}
+
+// loadedSkills returns the names of the skills already loaded in the session, whether pre-invoked or called by the model.
+func loadedSkills(s *Session) map[string]bool {
+	res := map[string]bool{}
+	for _, msg := range s.Messages(FilterByType(MessageTypeCall), FilterByTool(LoadSkillName)) {
+		content, err := s.UnmarshalMessageContent(msg)
+		if err != nil {
+			continue
+		}
+		if args, ok := content.(*LoadSkillArgs); ok {
+			res[args.Name] = true
+		}
+	}
+	return res
+}

runtime/ai/skill_references_test.go

@@ -0,0 +1,192 @@
+package ai_test
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1"
+	"github.com/rilldata/rill/runtime/ai"
+	"github.com/rilldata/rill/runtime/drivers"
+	"github.com/rilldata/rill/runtime/testruntime"
+	"github.com/stretchr/testify/require"
+)
+
+// turnFunc produces the assistant message the simulated model returns for one completion call.
+type turnFunc func(opts *drivers.CompleteOptions) *aiv1.CompletionMessage
+
+// scriptedAIService is a deterministic drivers.AIService for tests. Each Complete call consumes the next scripted
+// turn (falling back to a plain "done" reply once turns are exhausted) and records the messages it was given, so
+// tests can assert what the model saw.
+type scriptedAIService struct {
+	turns []turnFunc
+
+	mu     sync.Mutex
+	calls  int
+	inputs [][]*aiv1.CompletionMessage
+}
+
+var _ drivers.AIService = (*scriptedAIService)(nil)
+
+func (s *scriptedAIService) Complete(_ context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.inputs = append(s.inputs, opts.Messages)
+
+	var turn turnFunc
+	if s.calls < len(s.turns) {
+		turn = s.turns[s.calls]
+	}
+	s.calls++
+	if turn == nil {
+		turn = textTurn("done")
+	}
+
+	return &drivers.CompleteResult{
+		Message:      turn(opts),
+		Provider:     "scripted",
+		InputTokens:  1,
+		OutputTokens: 1,
+	}, nil
+}
+
+// textTurn makes the model reply with a plain text message (ending the tool loop).
+func textTurn(text string) turnFunc {
+	return func(_ *drivers.CompleteOptions) *aiv1.CompletionMessage {
+		return &aiv1.CompletionMessage{
+			Role:    "assistant",
+			Content: []*aiv1.ContentBlock{{BlockType: &aiv1.ContentBlock_Text{Text: text}}},
+		}
+	}
+}
+
+// newSkillReferencesSession creates a session on a project with analyst, developer and always-apply skills, backed by the given simulated model.
+func newSkillReferencesSession(t *testing.T, script *scriptedAIService) *ai.Session {
+	rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{
+		Files: map[string]string{
+			"rill.yaml": ``,
+			"skills/monthly-close/SKILL.md": `---
+description: Runs the monthly close analysis.
+agents: [analyst]
+---
+Compare revenue month over month.`,
+			"skills/churn-review/SKILL.md": `---
+description: Reviews customer churn.
+agents: [analyst, developer]
+---
+List the countries with the most churned customers.`,
+			"skills/glossary/SKILL.md": `---
+description: Business glossary.
+agents: [analyst]
+always_apply: true
+---
+ARPU excludes trial users.`,
+			// Without agents, a skill only applies to the developer agent.
+			"skills/dev-conventions/SKILL.md": `---
+description: Development conventions.
+---
+Name models in snake_case.`,
+		},
+	})
+	testruntime.RequireReconcileState(t, rt, instanceID, 5, 0, 0)
+
+	s := newSession(t, rt, instanceID)
+	s.SetLLM(func(_ context.Context) (drivers.AIService, func(), error) {
+		return script, func() {}, nil
+	})
+	return s
+}
+
+// loadedSkillNames returns the names of the skills loaded with load_skill as sub-calls of the given call, in order.
+func loadedSkillNames(s *ai.Session, callID string) []string {
+	var names []string
+	for _, call := range s.Messages(ai.FilterByParent(callID), ai.FilterByType(ai.MessageTypeCall), ai.FilterByTool(ai.LoadSkillName)) {
+		names = append(names, s.MustUnmarshalMessageContent(call).(*ai.LoadSkillArgs).Name)
+	}
+	return names
+}
+
+// loadSkillResults returns the tool result content of each load_skill call in the completion messages, by skill name.
+func loadSkillResults(messages []*aiv1.CompletionMessage) map[string]string {
+	names := map[string]string{} // tool call ID -> skill name
+	res := map[string]string{}
+	for _, m := range messages {
+		for _, block := range m.Content {
+			if call := block.GetToolCall(); call != nil && call.Name == ai.LoadSkillName {
+				names[call.Id] = call.Input.AsMap()["name"].(string)
+			}
+			if result := block.GetToolResult(); result != nil {
+				if name, ok := names[result.Id]; ok {
+					res[name] = result.Content
+				}
+			}
+		}
+	}
+	return res
+}
+
+// TestAnalystLoadsReferencedSkills verifies that the analyst loads the skills referenced with a chat-reference tag in the prompt before the model's first turn,
+// once per distinct analyst skill, and ignores references to skills that don't exist or don't apply to the analyst.
+func TestAnalystLoadsReferencedSkills(t *testing.T) {
+	script := &scriptedAIService{turns: []turnFunc{textTurn("done")}}
+	s := newSkillReferencesSession(t, script)
+
+	prompt := `<chat-reference>type="skill" skill="monthly-close"</chat-reference> for March, then ` +
+		`<chat-reference>skill="churn-review" type="skill"</chat-reference> and ` +
+		`<chat-reference>skill="monthly-close" type="skill"</chat-reference> again. ` +
+		`<chat-reference>type="skill" skill="does-not-exist"</chat-reference> ` +
+		`<chat-reference>type="skill" skill="dev-conventions"</chat-reference> ` +
+		`<chat-reference>type="skill" skill="glossary"</chat-reference> ` +
+		`<chat-reference>type="metricsView" metricsView="orders"</chat-reference>`
+	res, err := s.CallTool(t.Context(), ai.RoleUser, ai.AnalystAgentName, nil, &ai.AnalystAgentArgs{Prompt: prompt})
+	require.NoError(t, err)
+
+	// The always-apply glossary is pre-loaded once; the referenced analyst skills are loaded once each, in order of first reference.
+	require.Equal(t, []string{"glossary", "monthly-close", "churn-review"}, loadedSkillNames(s, res.Call.ID))
+
+	// The referenced skills' bodies were in the model's input on its first turn.
+	require.NotEmpty(t, script.inputs)
+	results := loadSkillResults(script.inputs[0])
+	require.Contains(t, results["monthly-close"], "Compare revenue month over month.")
+	require.Contains(t, results["churn-review"], "List the countries with the most churned customers.")
+	require.NotContains(t, results, "dev-conventions")
+	require.NotContains(t, results, "does-not-exist")
+}
+
+// TestAnalystSkipsSkillsAlreadyLoaded verifies that a skill whose body is already in the conversation is not loaded again when referenced in a later turn.
+func TestAnalystSkipsSkillsAlreadyLoaded(t *testing.T) {
+	script := &scriptedAIService{turns: []turnFunc{textTurn("first"), textTurn("second")}}
+	s := newSkillReferencesSession(t, script)
+
+	prompt := `<chat-reference>type="skill" skill="monthly-close"</chat-reference> for March`
+	res1, err := s.CallTool(t.Context(), ai.RoleUser, ai.AnalystAgentName, nil, &ai.AnalystAgentArgs{Prompt: prompt})
+	require.NoError(t, err)
+	require.Equal(t, []string{"glossary", "monthly-close"}, loadedSkillNames(s, res1.Call.ID))
+
+	// The same skill and the always-apply one, referenced again in a later turn.
+	prompt = `<chat-reference>type="skill" skill="monthly-close"</chat-reference> and ` +
+		`<chat-reference>type="skill" skill="glossary"</chat-reference> for April`
+	res2, err := s.CallTool(t.Context(), ai.RoleUser, ai.AnalystAgentName, nil, &ai.AnalystAgentArgs{Prompt: prompt})
+	require.NoError(t, err)
+	require.Empty(t, loadedSkillNames(s, res2.Call.ID))
+}
+
+// TestAnalystLoadsReferencedSkillsOnEveryTurn verifies that skills referenced in a later turn of the conversation are loaded in that turn.
+func TestAnalystLoadsReferencedSkillsOnEveryTurn(t *testing.T) {
+	script := &scriptedAIService{turns: []turnFunc{textTurn("first"), textTurn("second")}}
+	s := newSkillReferencesSession(t, script)
+
+	res1, err := s.CallTool(t.Context(), ai.RoleUser, ai.AnalystAgentName, nil, &ai.AnalystAgentArgs{Prompt: "Hello"})
+	require.NoError(t, err)
+	require.Equal(t, []string{"glossary"}, loadedSkillNames(s, res1.Call.ID))
+
+	prompt := `<chat-reference>type="skill" skill="monthly-close"</chat-reference> for March`
+	res2, err := s.CallTool(t.Context(), ai.RoleUser, ai.AnalystAgentName, nil, &ai.AnalystAgentArgs{Prompt: prompt})
+	require.NoError(t, err)
+	require.Equal(t, []string{"monthly-close"}, loadedSkillNames(s, res2.Call.ID))
+
+	// The body was in the model's input on the second turn's first completion.
+	require.Len(t, script.inputs, 2)
+	require.Contains(t, loadSkillResults(script.inputs[1])["monthly-close"], "Compare revenue month over month.")
+}

web-common/src/features/chat/core/context/config.ts

@@ -18,6 +18,7 @@ import {
   getLabelForComponent,
 } from "@rilldata/web-common/features/canvas/components/util.ts";
 import type { ChartSpec } from "@rilldata/web-common/features/components/charts/types.ts";
+import { SquareSlashIcon } from "lucide-svelte";
 
 type ContextConfigPerType = {
   editable: boolean;
@@ -140,4 +141,10 @@ export const InlineContextConfig: Record<
       `From ${InlineContextConfig[InlineContextType.Model].getLabel(ctx, meta)}`,
     getIcon: (ctx) => fieldTypeToSymbol(ctx.columnType ?? ""),
   },
+
+  [InlineContextType.Skill]: {
+    editable: false,
+    getLabel: (ctx) => ctx.skill ?? "",
+    getIcon: () => SquareSlashIcon,
+  },
 };