ai: Add a node embedding watcher #27204

hugoShaka · 2023-06-01T14:17:09Z

This PR adds the embedding core logic in auth:

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

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

xacrimon

first pass

api/proto/teleport/embedding/v1/embedding.proto

lib/ai/embedding.go

lib/auth/auth.go

lib/services/embeddings.go

lib/services/local/embeddings_test.go

lib/services/local/embeddings.go

lib/services/embeddings_test.go

lib/services/embeddings.go

api/proto/teleport/embedding/v1/embedding.proto

lib/ai/embedding.go

lib/auth/auth.go

lib/services/embeddings.go

jakule · 2023-06-08T17:03:56Z

lib/services/embeddings.go

+// It keeps tracks of which node has been embedded and which node requires embedding.
+// The embedding happens asynchronously as calling the openAI API by batch is
+// much quicker, stable and efficient.
+type nodeEmbeddingCollector struct {


This struct has 5 members, and 4 of them are responsible for synchronization. Any way to simplify it?

I don't think, 3 of them are here in all other collectors. The initialization channel and its fuse are here to please the resourceCollector interface. Same for the stale boolean.

The only synchronization structure used by our custom logic is the mutex.

lib/services/embeddings.go

jakule · 2023-06-08T17:08:38Z

lib/services/embeddings.go

+// embedding on all nodes needing embeddings. The embeddings are then inserted
+// into the vector index. This process is ran asynchronously to reduce the load
+// and leverage OpenAI's batch embedding API.
+func (n *nodeEmbeddingCollector) RunIndexation(ctx context.Context) error {


If I understand this correctly, if two goroutines run this method, both will read all nodes, then both will create embeddings and update the index with the same values. Shouldn't we just take the exclusive lock at the beginning and double-check if the embedding generation is needed to avoid it?

The only sensitive part in this is when the full sync runs and the function getResourcesAndUpdateCurrent acquires the lock.

There's a single routine triggering the embeddings so we should not be doing double work, but even if this happens this does not affect consistency because of the timestamp check.

lib/services/embeddings_test.go

- 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

jakule

I'm fine with the code as long as @justinas confirms the OpenAI API key logic.

jakule · 2023-06-12T19:53:09Z

lib/config/configuration.go

@@ -697,6 +697,16 @@ func applyAuthConfig(fc *FileConfig, cfg *servicecfg.Config) error {
 		cfg.Auth.Preference.SetDisconnectExpiredCert(fc.Auth.DisconnectExpiredCert.Value)
 	}

+	if fc.Auth.Assist != nil && fc.Auth.Assist.OpenAI != nil {


What will happen if the key is not provided in the Auth but in the Proxy? I think we have two options:

Disable the Assist as we cannot create embeddings

Make the embedding optional and let people use the feature without embeddings.

I personally vote for the first one and expend in the future if needed.

I think the auth endpoint doing the embedding lookup should return an error, and the proxies should handle it gracefully by bailing out and returning an error message saying embeddings must be enabled. I don't think it adds much value to handle cases where embeddings are disabled on the auth side, users will just have a poor UX.

lib/service/service.go

lib/service/servicecfg/auth.go

Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com>

lib/services/embeddings.go

lib/ai/embedding.go

lib/services/embeddings.go

lib/ai/embedding.go

api/proto/teleport/embedding/v1/embedding.proto

lib/ai/embedding.go

lib/services/local/embeddings.go

lib/services/embeddings.go

rosstimothy · 2023-06-14T13:18:26Z

lib/services/embeddings.go

+func (n *nodeEmbeddingCollector) NodeCount(needsEmbedding bool) int {
+	count := 0
+	n.mutex.Lock()
+	defer n.mutex.Unlock()
+	for _, node := range n.currentNodes {
+		if node.needsEmbedding == needsEmbedding {
+			count += 1
+		}
+	}
+	return count
+}


Do we need to calculate this each time the function is called? If the collector controls when things are added and when they are embedded shouldn't it be able to always know the count without iterating currentNodes?

We can optimize this function, but I think it would defeat this counter's purpose. Tests need a way to peak into the collector's state to know what is happening (and the tests are living in a separate services_test package, a private function is not usable here). Making this counter smarter and more efficient increases the risks of it returning an inaccurate value because of some collector bug.

rosstimothy · 2023-06-14T13:25:36Z

lib/services/embeddings.go

+// It keeps tracks of which node has been embedded and which node requires embedding.
+// The embedding happens asynchronously as calling the openAI API by batch is
+// much quicker, stable and efficient.
+type nodeEmbeddingCollector struct {


I think this collector is responsible for too much business logic. A lock watcher for example monitors and maintains the set of active locks and can be used to find a matching lock, but it does not do any connection termination as a result of a lock being active. I don't think any knowledge of how/when to perform an embedding should be given to the watcher.

Is this comment a blocker, or can we address it later when we'll extend embeddings to other resource kinds? This PR is somewhat important to unlock the rest of the Assist PRs for the next features and it would be easier for us to redesign this later.

I can extract the embedding logic out of the collector, but this will complexify even more an already complex locking situation as the collector and something else would need to sync to ensure embedding consistency. Do you think of a specific design that would allow consistent asynchronous embedding not triggered in the watcher?

I don't think that a resource watcher should know how to do anything other than monitor a resource and notify/make available the current state of that resource.

Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com>

GetEmbeddings now returns a stream

lib/services/local/embeddings.go

lib/services/embeddings.go

lib/services/local/embeddings.go

Co-authored-by: Edoardo Spadolini <edoardo.spadolini@goteleport.com>

…gs service

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.

public-teleport-github-review-bot · 2023-06-20T23:52:08Z

@hugoShaka - this PR will require admin approval to merge due to its size. Consider breaking it up into a series smaller changes.

r0mant

Bot.

* 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>

* [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>

hugoShaka changed the title ~~Add a node embedding watcher~~ ai: Add a node embedding watcher Jun 1, 2023

hugoShaka force-pushed the hugo/ai-embedding-watcher branch from 8fe4c68 to 7afeb85 Compare June 7, 2023 19:11

hugoShaka requested review from jakule and xacrimon June 8, 2023 12:42

xacrimon reviewed Jun 8, 2023

View reviewed changes

lib/services/local/embeddings_test.go Show resolved Hide resolved

lib/services/local/embeddings.go Show resolved Hide resolved

lib/services/embeddings_test.go Outdated Show resolved Hide resolved

lib/services/embeddings.go Show resolved Hide resolved

jakule reviewed Jun 8, 2023

View reviewed changes

hugoShaka force-pushed the hugo/ai-embedding-watcher branch 2 times, most recently from 37bcc11 to 5e9f6ff Compare June 9, 2023 14:08

hugoShaka requested review from jakule and xacrimon June 9, 2023 18:08

hugoShaka added 2 commits June 9, 2023 15:07

ai: adapt embeddings to the vector index

82de770

hugoShaka force-pushed the hugo/ai-embedding-watcher branch from ae99a34 to 82de770 Compare June 12, 2023 12:53

hugoShaka changed the base branch from dev-ai to master June 12, 2023 12:53

hugoShaka added 3 commits June 12, 2023 08:56

Merge branch 'master' into hugo/ai-embedding-watcher

be800cf

fixup! ai: adapt embeddings to the vector index

56f07f7

fixup! fixup! ai: adapt embeddings to the vector index

a45d571

hugoShaka marked this pull request as ready for review June 12, 2023 13:44

github-actions bot added the size/lg label Jun 12, 2023

github-actions bot requested review from ryanclark and tigrato June 12, 2023 13:45

jakule approved these changes Jun 12, 2023

View reviewed changes

Update lib/service/service.go

9a39a5f

Co-authored-by: Jakub Nyckowski <jakub.nyckowski@goteleport.com>

jakule reviewed Jun 12, 2023

View reviewed changes

lib/services/embeddings.go Outdated Show resolved Hide resolved

espadolini reviewed Jun 13, 2023

View reviewed changes

lib/ai/embedding.go Outdated Show resolved Hide resolved

xacrimon approved these changes Jun 13, 2023

View reviewed changes

lib/ai/embedding.go Outdated Show resolved Hide resolved

lib/services/embeddings.go Outdated Show resolved Hide resolved

jakule reviewed Jun 14, 2023

View reviewed changes

lib/ai/embedding.go Outdated Show resolved Hide resolved

espadolini reviewed Jun 14, 2023

View reviewed changes

lib/services/local/embeddings.go Outdated Show resolved Hide resolved

rosstimothy reviewed Jun 14, 2023

View reviewed changes

hugoShaka and others added 4 commits June 14, 2023 09:47

address feedback pt.1

a2103e2

address feedback pt.2: store protobuf message in backend

002bb8c

address feedback pt.3: have GetEmbeddings return a stream

8d3d703

Update lib/services/embeddings.go

e02c251

Co-authored-by: rosstimothy <39066650+rosstimothy@users.noreply.github.com>

hugoShaka requested a review from espadolini June 14, 2023 15:35

espadolini reviewed Jun 14, 2023

View reviewed changes

lib/services/local/embeddings.go Outdated Show resolved Hide resolved

lib/services/embeddings.go Outdated Show resolved Hide resolved

lib/services/embeddings.go Outdated Show resolved Hide resolved

lib/services/local/embeddings.go Show resolved Hide resolved

hugoShaka and others added 7 commits June 14, 2023 13:30

Apply suggestions from code review

04cf97a

Co-authored-by: Edoardo Spadolini <edoardo.spadolini@goteleport.com>

address feedback pt.4: extract embedding logic out of Embeddings service

15186dc

fixup! address feedback pt.4: extract embedding logic out of Embeddin…

98eab77

…gs service

address feedback pt.5: simpler error handling when embedding fails

f8c86cc

fix tests pt.1

6a7d87e

fix tests pt.2

5ec0661

fix tests pt.3

5cbd303

hugoShaka requested a review from rosstimothy June 14, 2023 19:08

[Assist] Replace embedding watcher (#27953)

3895105

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.

r0mant approved these changes Jun 21, 2023

View reviewed changes

public-teleport-github-review-bot bot removed request for tigrato and ryanclark June 21, 2023 00:46

jakule added this pull request to the merge queue Jun 21, 2023

Merged via the queue into master with commit c326b8c Jun 21, 2023
27 checks passed

jakule deleted the hugo/ai-embedding-watcher branch June 21, 2023 01:47

jakule restored the hugo/ai-embedding-watcher branch June 21, 2023 04:11

zmb3 mentioned this pull request Jul 20, 2023

tsh/tctl: remove lib/auth dependency #29375

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ai: Add a node embedding watcher #27204

ai: Add a node embedding watcher #27204

hugoShaka commented Jun 1, 2023 •

edited

xacrimon left a comment

jakule Jun 8, 2023

hugoShaka Jun 9, 2023

jakule Jun 8, 2023

hugoShaka Jun 9, 2023

jakule left a comment

jakule Jun 12, 2023

hugoShaka Jun 12, 2023 •

edited

rosstimothy Jun 14, 2023

hugoShaka Jun 14, 2023

rosstimothy Jun 14, 2023

hugoShaka Jun 14, 2023 •

edited

rosstimothy Jun 14, 2023

public-teleport-github-review-bot bot commented Jun 20, 2023

r0mant left a comment

ai: Add a node embedding watcher #27204

ai: Add a node embedding watcher #27204

Conversation

hugoShaka commented Jun 1, 2023 • edited

xacrimon 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

jakule left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hugoShaka Jun 12, 2023 • edited

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

hugoShaka Jun 14, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

public-teleport-github-review-bot bot commented Jun 20, 2023

r0mant left a comment

Choose a reason for hiding this comment

hugoShaka commented Jun 1, 2023 •

edited

hugoShaka Jun 12, 2023 •

edited

hugoShaka Jun 14, 2023 •

edited