assist: move embeddings in a separate library, add tool config, opportunistic node lookup

[hugoShaka]

This PR mitigates https://github.com/gravitational/teleport.e/issues/1820.

On the vast majority of Teleport Team clusters, the user only has access to a few nodes, and the cluster is small. In such cases, we can bypass the embedding lookup and feed the model with all allowed nodes directly from the proxy node watcher's cache. This allows users who just added a node to be able to use it with Assist. Assist behaviour will degrade to what we know today on larger clusters: with up to 20min of latency due to the embedding process.

This PR also adds 2 things:

• refactor the lib/ai module to extract embedding types into lib/ai/embedding. This was made to break a cyclic dependency between lib/services and lib/ai/model. (lib/ai/model needs accessChecker from lib/services lib/services needed Embedding from lib/ai, and lib/ai depends a lot on lib/ai/model)
• refactor the Agent to add a ToolsConfig structure allowing to configure tools and send all the clients they need while reducing the number of parameters needed to create a new Agent.

[jakule]

This looks a bit hacky. Now if you pass DisableEmbeddingsTool you also must pass username which won't be used.

[jakule]

The values here should not change. I think that because now the embedder is optional is disabled as a tool which is not correct. We still need to include the "get available node" tools in the prompt. We just need to use a different implementation for fetching the nodes.

[hugoShaka]

The test was previously instantiating a tool that would panic if invoked. As embedding tool was not used for the token counting test it did not crash.

I'm not sure what we want to test here except the tokens are counted, but IMO it is saner to remove broken tools from the test. If a test really depends on this tool, it should make it available and pass valid configuration.

[hugoShaka]

We'll have valid use-cases to enable/disable tools. e.g. Athena query building can be dropped if the user doesn't have access to audit logs. I think a modular agent toolbox is a pattern we'll benefit from, especially if we introduce new tools but want to reduce the token footprint.

[jakule]

I do understand that, but then you may end up with zero tools and very bad UX. I think that saving tokens + not confusing LLM are good arguments, but we should be smarter about when each tool is available. Can we for now still require the node embedded and fix the situation later (in a proper way)?

[hugoShaka]

I get what you're saying and have a proposal: would you be happy if, instead of having a hacky configuration that can enable/disable tools arbitrarily for the sake of testing, we had an interface with a method that returns the tools. This way we could have a standard implementation for real invocations, and tests could define their own custom tools. Like

type Toolbox interface {
        CheckAndSetDefaults()
        GetTools(context.Context) []model.Tool
}

type DefaultToolbox struct {
        EmbeddingsClient assist.AssistEmbeddingServiceClient
	AccessChecker services.AccessChecker
	NodeClient *services.NodeWatcher
}

// Returns [executionTool, embeddingTool] and all our future tools
func (tb *DefaultToolbox) GetTools(ctx context.Context) []model.Tool {...}

And in tests we could do

type dummyToolbox struct {
    testTools []model.Tools
}

// Returns what dummy or real tools contained in the toolbox
func (tb *dummyToolbox) GetTools(ctx context.Context) []model.Tool

This would allow tests to safely choose which tools are here, and mock tools or use real tools when needed, while keeping the "real" toolbox safe and always containing all tools.

[jakule]

I like the idea of separating the tools from the agent itself. I think we will need that anyways, as when I was working on the "in-context SSH Assist" I needed a different agent with a completely different set of tools.
I think we could even rename the DefaultToolbox to ChatToolbox as it contains all tools required by the chat. Then as you mentioned, other agents (like tests) can implement a different set of available tools.
@xacrimon It this change going to conflict with your access request changes? If yes, we can ship it after this and your PR are merged.

[xacrimon]

Yep, im refactoring that area a little, sounds good.

[jakule]

As we talked at the meeting. The code overall looks good. The agent logic needs to change anyways to accommodate other Assisst-related changes. It can be done in a separate PR.