Build an AI assistant

// newHelpdesk returns the agent that answers questions about Tidy from
// the help center, with two tools over the articles: a hybrid search
// (their meaning and their words), and reading one.
func newHelpdesk(articles *ai.Embeddings[Article]) ai.Agent {
	return ai.Agent{
		Name: "helpdesk",
		Instructions: "You answer questions about Tidy, a to-do app, from its help-center articles. " +
			"Search the articles, read the best match if its passage isn't enough, and answer briefly, naming the article. " +
			"If the articles don't say, say you don't know.",
		Tools: []ai.Tool{
			articles.Tool("search_articles", "Search the help center by meaning and words; returns the best articles' IDs, titles and passages", 3),
			readArticle,
		},
		MaxSteps: 6,
	}
}

// ReadInput is what the model sends to read_article.
type ReadInput struct {
	ID int64 `json:"id" description:"The article's ID, from search_articles" validate:"required"`
}

var readArticle = ai.Func("read_article", "Read a help-center article",
	func(ctx context.Context, in ReadInput) (Article, error) {
		return db.Find[Article](ctx, in.ID) // not found: 404, told to the model
	})

client, err := ai.ForApp(app, anthropic.Driver(), openai.Driver(), openai.CompatibleDriver())
if err != nil {
	return nil, err
}
client.TrackUsage(ai.UsageConfig{ // the ai_usage table, and budgets
	Budget: func(context.Context, string) (ai.Budget, error) {
		return ai.Budget{Tokens: settings.DailyTokens, Per: 24 * time.Hour}, nil
	},
	// Prices: map[string]ai.Price{"model-name": {Input: …, Output: …}}, // per million tokens
})
articles, err := ai.EmbeddingsFor(app, articleEmbeddings) // articles_embeddings; ai:embed
if err != nil {
	return nil, err
}
anetos.Provide(app, articles)
helpdesk := newHelpdesk(articles)
if err := ai.QueueAgents(app, helpdesk); err != nil { // answers in the background
	return nil, err
}

// Start stores a new conversation with the question, and shows it: the
// page streams the answer.
func (Handlers) Start(c *web.Ctx, in NewChat) (web.Responder, error) {
	userID, err := auth.CurrentID(c)
	if err != nil {
		return nil, err
	}
	conv, err := ai.StartConversation(c, userID, truncate(in.Prompt, 60))
	if err != nil {
		return nil, err
	}
	if err := conv.Add(c, ai.UserMessage(in.Prompt)); err != nil {
		return nil, err
	}
	return web.Redirect(fmt.Sprintf("/chat/%d", conv.ID)), nil
}

// Send stores the question and returns it, with a place for the answer
// that streams it from Reply.
func (Handlers) Send(c *web.Ctx, in Question) (web.Responder, error) {
	conv, err := conversation(c, in.ID)
	if err != nil {
		return nil, err
	}
	if err := conv.Add(c, ai.UserMessage(in.Prompt)); err != nil {
		return nil, err
	}
	return page("exchange", map[string]any{"Conversation": conv, "Prompt": in.Prompt}), nil
}

// Reply streams the answer to the conversation's last question, and
// stores it.
func (h Handlers) Reply(c *web.Ctx, in ChatPath) (web.Responder, error) {
	conv, err := conversation(c, in.ID)
	if err != nil {
		return nil, err
	}
	msgs, err := conv.Messages(c)
	if err != nil {
		return nil, err
	}
	answer := conv.StreamReply(c, h.helpdesk)
	if len(msgs) == 0 || msgs[len(msgs)-1].Role != ai.RoleUser || conv.Status == ai.StatusQueued {
		answer = func(func(ai.Event, error) bool) {} // answered already: a browser that reconnected
	}
	return web.ResponderFunc(func(c *web.Ctx) error { return ai.SSE(c, answer) }), nil
}

<!-- illustrative -->
<script src="{{call .Asset "htmx.min.js"}}"></script>
<script src="{{call .Asset "htmx-ext-sse.min.js"}}"></script>

<div class="msg assistant" hx-ext="sse" sse-connect="/chat/7/reply" sse-close="done">
  <p class="tool" sse-swap="tool"></p>
  <p sse-swap="text" hx-swap="beforeend"></p>
  <p class="error" sse-swap="error"></p>
</div>

// Later stores the question for a queue job to answer, and returns a
// placeholder that polls Status.
func (h Handlers) Later(c *web.Ctx, in Question) (web.Responder, error) {
	conv, err := conversation(c, in.ID)
	if err != nil {
		return nil, err
	}
	if err := conv.Add(c, ai.UserMessage(in.Prompt)); err != nil {
		return nil, err
	}
	if err := conv.QueueReply(c, h.helpdesk); err != nil {
		return nil, err
	}
	return page("queued", map[string]any{"Conversation": conv, "Prompt": in.Prompt}), nil
}

// Status returns the queued answer when it's there, else the placeholder
// again. If the job left the question unanswered (a question asked in
// the meantime), the answer streams instead.
func (Handlers) Status(c *web.Ctx, in ChatPath) (web.Responder, error) {
	conv, err := conversation(c, in.ID)
	if err != nil {
		return nil, err
	}
	switch conv.Status {
	case ai.StatusQueued:
		return page("waiting", map[string]any{"Conversation": conv}), nil
	case ai.StatusFailed:
		return page("failed", map[string]any{"Error": conv.Error}), nil
	}
	msgs, err := conv.Messages(c)
	if err != nil {
		return nil, err
	}
	if n := len(msgs); n > 0 && msgs[n-1].Role == ai.RoleUser {
		return page("answer", map[string]any{"Conversation": conv}), nil
	}
	all := bubbles(msgs)
	return page("bubbles", all[max(len(all)-1, 0):]), nil
}

today := anetos.Now(c).UTC().Truncate(24 * time.Hour)
usage, err := ai.TotalUsage(c, userID, today)
if err != nil {
	return err
}

func TestAssistant(t *testing.T) {
	// The model's replies, scripted: search, read, answer.
	app := anetostest.New(t, setup, anetostest.FakeAI(
		ai.FakeToolCall("search_articles", map[string]string{"query": "How do I export my lists?"}),
		ai.FakeToolCall("read_article", ReadInput{ID: 1}),
		ai.FakeText("Open Settings, then Data, and choose Export (Export your lists)."),
	))
	addArticles(t, app)
	signIn(t, app, "Ada")

	app.PostForm("/chat", url.Values{"prompt": {"How do I export my lists?"}}).Follow().
		AssertSee("How do I export my lists?", `sse-connect="/chat/1/reply"`)
	// The answer streams as server-sent events, and is stored.
	app.Get("/chat/1/reply").AssertOK().
		AssertSee("event: tool\ndata: search_articles", "event: tool\ndata: read_article",
			"event: text\ndata: Open ", "event: done")
	app.Get("/chat/1").AssertSee("Used: search_articles, read_article", "Open Settings, then Data, and choose Export").
		AssertDontSee("sse-connect")

	// The tools ran for real: the search found the article, with its
	// passage, which the model read.
	reqs := app.AI().Requests()
	if got := reqs[1].Messages[2].Parts[0].(ai.ToolResult).Content; !strings.HasPrefix(got, `[{"id":1,"title":"Export your lists","text":"Export your lists\n\nOpen Settings`) {
		t.Errorf("search results: %s", got)
	}
	app.AssertPrompted(func(r ai.Request) bool { return strings.Contains(r.System, "help-center articles") })

	// A browser that reconnects gets no second answer.
	app.Get("/chat/1/reply").AssertOK().AssertDontSee("event: text")
}

func TestAnswerLater(t *testing.T) {
	// QUEUE_DRIVER is sync by default: the job runs at once.
	app := anetostest.New(t, setup, anetostest.FakeAI(ai.FakeText("Hello!"), ai.FakeText("Shared lists need a Team plan.")))
	signIn(t, app, "Ada")
	app.PostForm("/chat", url.Values{"prompt": {"Hi"}})
	app.Get("/chat/1/reply")

	app.PostForm("/chat/1/later", url.Values{"prompt": {"Can I share a list?"}}).
		AssertSee("Can I share a list?", `hx-get="/chat/1/status"`)
	app.Get("/chat/1/status").AssertSee("Shared lists need a Team plan.").AssertDontSee("Working on it")
}