[v14] Add the ability to run a specific tool to Assist.#33640
[v14] Add the ability to run a specific tool to Assist.#33640jakule merged 2 commits intobranch/v14from
Conversation
This commit contains: - a refactoring of the model action logic to be exposed through the `DoAction` function - Expose `RunTool()` capability through ai.Client and assist.Assist - Add the "audit-query" action that invokes directly the audit generation tool
github-actions
Bot
requested review from
hugoShaka,
rosstimothy,
strideynet and
zmb3
| Reasoning: "Tool invoked directly", | ||
| } | ||
|
|
||
| return agent.DoAction(ctx, client.svc, action) |
There was a problem hiding this comment.
nit: We should wrap the error. I'm not sure what the other returns are from the context here.
| return agent.DoAction(ctx, client.svc, action) | |
| ?, ?, err = agent.DoAction(ctx, client.svc, action) | |
| return ?, ?, trace.Wrap(err) |
There was a problem hiding this comment.
I think we should, but I also don't think we should do that in a backport. We can address those comments in a separate PR targeting master.
| response, _, err := client.RunTool(ctx, toolCtx, tools.AuditQueryGenerationToolName, "List users who connected to a server as root") | ||
| require.NoError(t, err) | ||
| message, ok := response.(*output.StreamingMessage) | ||
| require.True(t, ok) |
There was a problem hiding this comment.
Suggestion: add more context so that we get more than "should be true" if the assertion fails
| require.True(t, ok) | |
| require.True(t, ok, "expected response to be an output.StreamingMessage, got %T, response) |
There was a problem hiding this comment.
I think we could change, but I also don't think we should do that in a backport. We can address those comments in a separate PR targeting master.
| response, _, err := client.RunTool(ctx, toolCtx, "Nodes names and labels retrieval", string(inputText)) | ||
| require.NoError(t, err) | ||
| message, ok := response.(*output.Message) | ||
| require.True(t, ok) |
There was a problem hiding this comment.
Suggestion: add more context so that we get more than "should be true" if the assertion fails
| require.True(t, ok) | |
| require.True(t, ok, "expected response to be an output.StreamingMessage, got %T, response |
There was a problem hiding this comment.
I think we could change, but I also don't think we should do that in a backport. We can address those comments in a separate PR targeting master.
| "github.com/gravitational/teleport/lib/ai/tokens" | ||
| ) | ||
|
|
||
| const AuditQueryGenerationToolName = "Audit Query Generation" |
There was a problem hiding this comment.
nit: missing docs on exported constant
There was a problem hiding this comment.
As above. It can be addressed in a separate PR.
|
This also does not look like a backport of #31112 |
* stream audit query responses * fixup! stream audit query responses * fixup! fixup! stream audit query responses
I fixed that, thx |
smallinsky
left a comment
smallinsky
left a comment
There was a problem hiding this comment.
LGTM
Tim left some valid comments but agree with Jakub that we should align it in master first and backport.
public-teleport-github-review-bot
Bot
removed request for
hugoShaka,
strideynet and
zmb3
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
5 participants
Backport of #31113