|
1 | 1 | package coderd
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "context" |
4 | 5 | "database/sql"
|
5 | 6 | "errors"
|
6 | 7 | "fmt"
|
| 8 | + "io" |
7 | 9 | "net/http"
|
| 10 | + "os" |
8 | 11 | "slices"
|
9 | 12 | "strings"
|
10 | 13 |
|
| 14 | + "github.com/anthropics/anthropic-sdk-go" |
| 15 | + anthropicoption "github.com/anthropics/anthropic-sdk-go/option" |
11 | 16 | "github.com/google/uuid"
|
12 | 17 |
|
| 18 | + "github.com/coder/aisdk-go" |
13 | 19 | "github.com/coder/coder/v2/coderd/audit"
|
14 | 20 | "github.com/coder/coder/v2/coderd/database"
|
15 | 21 | "github.com/coder/coder/v2/coderd/httpapi"
|
@@ -69,6 +75,69 @@ func (api *API) aiTasksPrompts(rw http.ResponseWriter, r *http.Request) {
|
69 | 75 | })
|
70 | 76 | }
|
71 | 77 |
|
| 78 | +func (api *API) generateTaskName(ctx context.Context, prompt, fallback string) (string, error) { |
| 79 | + // TODO(DanielleMaywood): |
| 80 | + // Should we extract this out into our typical coderd option handling? |
| 81 | + anthropicAPIKey := os.Getenv("ANTHROPIC_API_KEY") |
| 82 | + |
| 83 | + // The deployment doesn't have a valid external cloud AI provider, so we'll |
| 84 | + // fallback to the user supplied name for now. |
| 85 | + if anthropicAPIKey == "" { |
| 86 | + return fallback, nil |
| 87 | + } |
| 88 | + |
| 89 | + anthropicClient := anthropic.NewClient(anthropicoption.WithAPIKey(anthropicAPIKey)) |
| 90 | + |
| 91 | + messages, system, err := aisdk.MessagesToAnthropic([]aisdk.Message{ |
| 92 | + { |
| 93 | + Role: "system", |
| 94 | + Parts: []aisdk.Part{{ |
| 95 | + Type: aisdk.PartTypeText, |
| 96 | + Text: ` |
| 97 | + You are a task summarizer. |
| 98 | + You summarize AI prompts into workspace names. |
| 99 | + You will only respond with a workspace name. |
| 100 | + The workspace name **MUST** follow this regex /^[a-z0-9]+(?:-[a-z0-9]+)*$/ |
| 101 | + The workspace name **MUST** be 32 characters or **LESS**. |
| 102 | + The workspace name **MUST** be all lower case. |
| 103 | + The workspace name **MUST** end in a number between 0 and 100. |
| 104 | + The workspace name **MUST** be prefixed with "task". |
| 105 | + `, |
| 106 | + }}, |
| 107 | + }, |
| 108 | + { |
| 109 | + Role: "user", |
| 110 | + Parts: []aisdk.Part{{ |
| 111 | + Type: aisdk.PartTypeText, |
| 112 | + Text: prompt, |
| 113 | + }}, |
| 114 | + }, |
| 115 | + }) |
| 116 | + if err != nil { |
| 117 | + return "", err |
| 118 | + } |
| 119 | + |
| 120 | + stream := aisdk.AnthropicToDataStream(anthropicClient.Messages.NewStreaming(ctx, anthropic.MessageNewParams{ |
| 121 | + Model: anthropic.ModelClaude3_5HaikuLatest, |
| 122 | + Messages: messages, |
| 123 | + System: system, |
| 124 | + MaxTokens: 24, |
| 125 | + })) |
| 126 | + |
| 127 | + var acc aisdk.DataStreamAccumulator |
| 128 | + stream = stream.WithAccumulator(&acc) |
| 129 | + |
| 130 | + if err := stream.Pipe(io.Discard); err != nil { |
| 131 | + return "", err |
| 132 | + } |
| 133 | + |
| 134 | + if len(acc.Messages()) == 0 { |
| 135 | + return fallback, nil |
| 136 | + } |
| 137 | + |
| 138 | + return acc.Messages()[0].Content, nil |
| 139 | +} |
| 140 | + |
72 | 141 | // This endpoint is experimental and not guaranteed to be stable, so we're not
|
73 | 142 | // generating public-facing documentation for it.
|
74 | 143 | func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
|
@@ -104,8 +173,21 @@ func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
|
104 | 173 | return
|
105 | 174 | }
|
106 | 175 |
|
| 176 | + taskName, err := api.generateTaskName(ctx, req.Prompt, req.Name) |
| 177 | + if err != nil { |
| 178 | + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ |
| 179 | + Message: "Internal error generating name for task.", |
| 180 | + Detail: err.Error(), |
| 181 | + }) |
| 182 | + return |
| 183 | + } |
| 184 | + |
| 185 | + if taskName == "" { |
| 186 | + taskName = req.Name |
| 187 | + } |
| 188 | + |
107 | 189 | createReq := codersdk.CreateWorkspaceRequest{
|
108 |
| - Name: req.Name, |
| 190 | + Name: taskName, |
109 | 191 | TemplateVersionID: req.TemplateVersionID,
|
110 | 192 | TemplateVersionPresetID: req.TemplateVersionPresetID,
|
111 | 193 | RichParameterValues: []codersdk.WorkspaceBuildParameter{
|
|
0 commit comments