[Assist] Scaffold the chat-loop onto a multi-step thinking model #27075

xacrimon · 2023-05-29T11:28:23Z

This PR refactors our use the of the LLM and how we view it. This PR swaps the current single-shot thinking for a design where another message is used to take a decision on the next course of action. This is a necessary prerequisite for future additions.

This PR temporarily disables streaming as refactoring streaming would make this PR far too large. It will be re-added in a later PR.

Fixes https://github.com/gravitational/teleport.e/issues/1463

jakule

Can we add some tests? We already have examples of how to mock the OpenAI API.

lib/ai/model/agent.go

jakule · 2023-06-02T15:43:31Z

lib/ai/model/agent.go

+const (
+	actionFinalAnswer = "Final Answer"
+	actionException   = "_Exception"
+	maxIterations     = 15
+	maxElapsedTime    = 5 * time.Minute
+)
+
+var AssistAgent = &Agent{
+	tools: []Tool{
+		&commandExecutionTool{},
+	},
+}
+
+type Agent struct {
+	tools []Tool
+}
+
+type AgentAction struct {
+	action string
+	input  string
+	log    string
+}
+
+type AgentFinish struct {
+	// output must be Message or CompletionCommand
+	output any


Can we get GoDocs?

lib/ai/model/agent.go

lib/ai/model/error.go

lib/ai/model/tool.go

lib/assist/assist.go

flyinghermit

Few comments and questions but looks good overall.

flyinghermit · 2023-06-06T12:53:42Z

lib/ai/chat.go


-const maxResponseTokens = 2000


Any reason to remove this? FWIR, this was used to deal with OpenAI rate limits? (pardon if I am out of context)

This was used to avoid sometimes excessively large responses. With the new agent model and prompt setup, output size is better controlled and this only prevents the occasional large useful response.

lib/ai/model/agent.go

flyinghermit · 2023-06-06T13:12:57Z

lib/ai/model/agent.go

+	return thoughts
+}
+
+func parseJSONFromModel[T any](text string) (T, *invalidOutputError) {


Noted that JSON object will be generic. But I was also curious to see what type of sanitization we are dealing with here. Maybe a sample JSON reference would be helpful.

flyinghermit · 2023-06-06T13:39:22Z

lib/ai/model/prompt.go

+	return fmt.Sprintf(`You are Teleport, a tool that users can use to connect to Linux servers and run relevant commands, as well as have a conversation.
+A Teleport cluster is a connectivity layer that allows access to a set of servers. Servers may also be referred to as nodes.
+Nodes sometimes have labels such as "production" and "staging" assigned to them. Labels are used to group nodes together.
+You will engage in professional conversation with the user and help accomplish tasks such as executing tasks


tasks such as executing tasks

Should we have a short sentence to explain "tasks"? Other details looks great 👍

I didn't find this improved the model in my testing so I opted to leave as is until we can spend more time on this.

flyinghermit · 2023-06-06T13:44:38Z

lib/ai/model/prompt.go

+
+const InitialAIResponse = `Hey, I'm Teleport - a powerful tool that can assist you in managing your Teleport cluster via OpenAI GPT-4.`
+
+func PromptCharacter(username string) string {


Does this prompts affect production chats? If yes, have we done A/B tests to check it. Not a blocker for this PR but something that would be good to chase for future.

It does, but this is also very similar to the old prompt with an added sentence or so. I have done significant testing on my end to make sure it behaves nicely.

lib/ai/model/agent.go

lib/ai/model/tool.go

… the llm

This test is removed altogether in #27075, but that has not been backported yet.

* Return errors over Assist WS (#27174) * Return errors over Assist WS * Add test * Address code review comments Add check for the close message in assist WS * Use UTC for getting time * Gracefully handle errors in Assist frontend (#27669) * Bump e to include teleport.e#1634 * Fix error assertions in the test This test is removed altogether in #27075, but that has not been backported yet. --------- Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com>

The tests were removed as a part of #27075. This PR updates the tests to use the new logic.

* Restore `lib/ai` tests The tests were removed as a part of #27075. This PR updates the tests to use the new logic. * Fix tests * Restore lib/web tests * GCI * Move test handler to a common place * Fix used token test * Add comment

) * agent scaffold conversion * command input validation * rename Agent.Think and replace debug logs with trace logs * doc * action docs * godocs * clarify * remove unused code * remove tests which relied on the old non-agent model interaction with the llm * fix broken e

* Restore `lib/ai` tests The tests were removed as a part of #27075. This PR updates the tests to use the new logic. * Fix tests * Restore lib/web tests * GCI * Move test handler to a common place * Fix used token test * Add comment

* [Assist] Scaffold the chat-loop onto a multi-step thinking model (#27075) * agent scaffold conversion * command input validation * rename Agent.Think and replace debug logs with trace logs * doc * action docs * godocs * clarify * remove unused code * remove tests which relied on the old non-agent model interaction with the llm * fix broken e * Add node name to the Assist execution result (#27635) * Add node name to the Assist execution result Currently, only node ID is returned on the command execution result in Assist. For better UX we want to display Node name which id more human friendly rather than a node ID which is a UUID. Adding the value to returned payload sounds cheaper than calling an API to get node names. * Add test * Extract commandExecResult struct * Fix test after rebase * Fix command execution test flakiness (#27704) Fix ``` --- FAIL: TestExecuteCommand (1.46s) testing.go:1206: TempDir RemoveAll cleanup: unlinkat /tmp/TestExecuteCommand3553793052/002/log/upload/streaming/default: directory not empty FAIL ``` error * [Assist] Fix panic when writing to one WS from multiple threads (#27828) * [Assist] Fix panic when writing to one WS from multiple threads Fixes gravitational/teleport.e#1650 * Remove mutex on SetReadDeadline * Move SetPongHandler * Fix typos * Fix command output showing when running on multiple nodes (#27936) * ai: Add a node embedding watcher (#27204) * ai: add embeddings basic support - add Embeddings service and its local implementation - add Embedding type and proto message - add nodeEmbeddingCollector tracking nodes - add NodeEmbeddingWatcher watching for events adn sending them to the collector - add the Embedder interface and its openai implementation * ai: adapt embeddings to the vector index * fixup! ai: adapt embeddings to the vector index * fixup! fixup! ai: adapt embeddings to the vector index * Update lib/service/service.go Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com> * address feedback pt.1 * address feedback pt.2: store protobuf message in backend * address feedback pt.3: have GetEmbeddings return a stream * Update lib/services/embeddings.go Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Edoardo Spadolini <edoardo.spadolini@goteleport.com> * address feedback pt.4: extract embedding logic out of Embeddings service * fixup! address feedback pt.4: extract embedding logic out of Embeddings service * address feedback pt.5: simpler error handling when embedding fails * fix tests pt.1 * fix tests pt.2 * fix tests pt.3 * [Assist] Replace embedding watcher (#27953) Change the way how the embeddings are calculated. Instead of creating a watcher in Auth, we will process all nodes every hour and process embeddings if any embeddings are missing or any node has been updated. --------- Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com> Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> Co-authored-by: Edoardo Spadolini <edoardo.spadolini@goteleport.com> * Restore `lib/ai` tests (#28077) * Restore `lib/ai` tests The tests were removed as a part of #27075. This PR updates the tests to use the new logic. * Fix tests * Restore lib/web tests * GCI * Move test handler to a common place * Fix used token test * Add comment * Remove duplicate imports (#27886) * [Assist] Remove the empty assist message (#28125) * [Assist] Remove the empty assist message Assist shows an empty message at the beginning of each conversation when reading it from DB. This PR fixes that behavior and adds a test to prevent this from happening in the future. * Address code review comments * Address code review comments * Skip embedding processor on Cloud Non-Team plan (#28197) * ai: compute opportinistic summary of command execution (#28033) * ai: compute opportinistic summary of command execution * ai: add streaming summary back after rebase on new front-end * Lint and fix tests pt.1 * reference nodes by name and add tests * Lint, fix tests and address feedback * Attempt to tame the stream close monster * fixup! Attempt to tame the stream close monster * [Assist] Do not close the WS after command execution (#28246) * Revert "fixup! Attempt to tame the stream close monster" This reverts commit 8537aa2. * Revert "Attempt to tame the stream close monster" This reverts commit e0c861d. * Do not close the WS after command execution * Fix tests and lint * fixup! Fix tests and lint * undo put web test command into constant --------- Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com> * [Assist] Include embeddings in the prompt (#28116) * [Assist] Include embeddings in the prompt * Add comments GCI Minor fixes * Move stuff * Fix tests * Fix tests * Fixes after rebase Apply code review suggestions. * Address review comments * After rebase fix * Improve error handling and embedding prompts; fix typos (#28403) * "Improve error handling and embedding prompts; fix typos" This commit encompasses several changes. First, an error handling routine has been added in AssistContext.tsx to properly close a WebSocket connection and finish all results. The intent is to ensure that execution fails gracefully when a session doesn't end normally. In tool.go, user instructions have been made more explicit to ensure users check access to nodes before generating any commands. It warns them that not checking access will cause error. Also, some minor typos were corrected in agent.go and messages.go for better readability. * "Refactor 'hosts' to 'nodes' in AI Tool Descriptions" This commit refactors the language from 'host' terminology to 'node' terminology in the AI tool's generated responses as the LLM seems to be confused when generating queries with embeddings. * Update expected test values in chat_test.go The expected values in three different tests in chat_test.go have been updated. This change was required because the underlying algorithm has been adjusted and these modifications will keep the tests aligned with the current algorithm's behavior. * Add missing imports * Introduce user preferences (#28291) * Add user preferences feature * Add missing license header * Fix the order of arguments to require.Equal * Update lib/web/userpreferences.go Co-authored-by: Michelle Bergquist <11967646+michellescripts@users.noreply.github.com> * Add a `GetUserPreferencesResponse` message * Remove unused logger * Use .Put instead of .Create/.Update * Add missing godoc * trace.Wrap the happy path --------- Co-authored-by: Michelle Bergquist <11967646+michellescripts@users.noreply.github.com> * Shut down embedding processor on graceful exit (#28356) * Refactor websocket termination and stream handling (#28452) * Refactor websocket termination and stream handling Refactored websocket stream shutdown and error handling. Replaced `Close()` with `SendCloseMessage()` for better control over the websocket connection termination process. Added checks for the validity of channels to prevent reading from closed channels. The commit also includes minor typo fixes. * Remove unused completedC * Remove unnecessary select blocks in terminal.go The select blocks used in terminal.go for reading data from channels were unnecessary as we were just pulling from a single channel. Removed the select block and directly attempted to read from the channel. These changes increase code readability and integrity by removing unnecessary select blocks. In the command_test.go, an explanatory comment was added for clarity. * Remove commented code * Replace trace.NewAggregate with trace.Wrap as aggregation is not needed. * Add the UI for Assist's settings (#28413) * Add the UI for Assist's settings * Add typing * Fix test by wrapping render in LayoutContextProvider * Run prettier * Assist: fix summary logic (#28487) * Update command.go * simplify export signature * assist: add classification code (#28221) * [Assist] Provide interactive updates during agent execution (#27893) * send progress update messages during agent thoughts * handle new output format * define json tags for serialized fields * use streaming api * fan streaming from model loop * fix streaming * stream progress updates * Update lib/assist/assist.go Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com> * remove useless mute * nits * Update lib/ai/model/agent.go Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com> * fix merge * fix misc * more misc fixes * what * what2 * weird eof errors? * Fix tests UI integration * Fix other tests * Linter fixes * Comment out token counting for assist streams to avoid race condition. * Fix more tests --------- Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com> * Remove console.log in AssistContext (#28607) --------- Co-authored-by: Joel <jwejdenstal@goteleport.com> Co-authored-by: Ryan Clark <ryan.clark@goteleport.com> Co-authored-by: Hugo Shaka <hugo.hervieux@goteleport.com> Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com> Co-authored-by: Edoardo Spadolini <edoardo.spadolini@goteleport.com> Co-authored-by: Zac Bergquist <zac.bergquist@goteleport.com> Co-authored-by: Justinas Stankevičius <justinas@users.noreply.github.com> Co-authored-by: Michelle Bergquist <11967646+michellescripts@users.noreply.github.com>

xacrimon changed the title ~~[Assist] Scaffold the model onto a multi-step thinking model~~ [Assist] Scaffold the chat-loop onto a multi-step thinking model May 29, 2023

xacrimon marked this pull request as ready for review June 2, 2023 13:28

xacrimon force-pushed the joel/agent-scaffold branch from cde7f63 to 66a3ff9 Compare June 2, 2023 13:28

github-actions bot requested review from ryanclark and flyinghermit June 2, 2023 13:29

github-actions bot added the size/sm label Jun 2, 2023

jakule reviewed Jun 2, 2023

View reviewed changes

flyinghermit approved these changes Jun 6, 2023

View reviewed changes

jakule approved these changes Jun 6, 2023

View reviewed changes

lib/ai/model/agent.go Show resolved Hide resolved

lib/ai/model/agent.go Show resolved Hide resolved

lib/ai/model/tool.go Show resolved Hide resolved

public-teleport-github-review-bot bot removed the request for review from ryanclark June 6, 2023 21:27

xacrimon added 9 commits June 8, 2023 15:47

agent scaffold conversion

5c1b56c

command input validation

27c3d1f

rename Agent.Think and replace debug logs with trace logs

469aa83

doc

ceb9ce5

action docs

622170d

godocs

acb3ae5

clarify

5d0e60f

remove unused code

5e79fdd

remove tests which relied on the old non-agent model interaction with…

9aa2144

… the llm

xacrimon force-pushed the joel/agent-scaffold branch from d056b66 to 9aa2144 Compare June 8, 2023 13:55

xacrimon enabled auto-merge June 8, 2023 13:55

fix broken e

985d136

xacrimon added this pull request to the merge queue Jun 8, 2023

Merged via the queue into master with commit 36dcc6b Jun 8, 2023
20 checks passed

xacrimon deleted the joel/agent-scaffold branch June 8, 2023 14:33

justinas added a commit that referenced this pull request Jun 16, 2023

Fix error assertions in the test

6c4518c

This test is removed altogether in #27075, but that has not been backported yet.

jakule added a commit that referenced this pull request Jun 20, 2023

Restore lib/ai tests

491bff5

The tests were removed as a part of #27075. This PR updates the tests to use the new logic.

jakule mentioned this pull request Jun 20, 2023

Restore lib/ai tests #28077

Merged

fheinecke mentioned this pull request Sep 26, 2023

Release 14.0.1 #32611

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Assist] Scaffold the chat-loop onto a multi-step thinking model #27075

[Assist] Scaffold the chat-loop onto a multi-step thinking model #27075

xacrimon commented May 29, 2023 •

edited

jakule left a comment

jakule Jun 2, 2023

flyinghermit left a comment

flyinghermit Jun 6, 2023

xacrimon Jun 8, 2023

flyinghermit Jun 6, 2023

flyinghermit Jun 6, 2023

xacrimon Jun 8, 2023

flyinghermit Jun 6, 2023

xacrimon Jun 8, 2023


		const InitialAIResponse = `Hey, I'm Teleport - a powerful tool that can assist you in managing your Teleport cluster via OpenAI GPT-4.`

		func PromptCharacter(username string) string {

[Assist] Scaffold the chat-loop onto a multi-step thinking model #27075

[Assist] Scaffold the chat-loop onto a multi-step thinking model #27075

Conversation

xacrimon commented May 29, 2023 • edited

jakule left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flyinghermit left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xacrimon commented May 29, 2023 •

edited