Skip to content

Commit 0fd0daf

Browse files
sawkasawka
authored
Reimplement wsh ai, fix text file attaching format (#2435)
1 parent 0e8eb83 commit 0fd0daf

File tree

17 files changed

+581
-229
lines changed

17 files changed

+581
-229
lines changed

‎cmd/wsh/cmd/wshcmd-ai.go‎

Lines changed: 122 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -4,166 +4,189 @@
44
package cmd
55

66
import (
7+
"encoding/base64"
78
"fmt"
89
"io"
10+
"net/http"
911
"os"
12+
"path/filepath"
1013
"strings"
1114

1215
"github.com/spf13/cobra"
13-
"github.com/wavetermdev/waveterm/pkg/waveobj"
16+
"github.com/wavetermdev/waveterm/pkg/util/utilfn"
1417
"github.com/wavetermdev/waveterm/pkg/wshrpc"
1518
"github.com/wavetermdev/waveterm/pkg/wshrpc/wshclient"
1619
"github.com/wavetermdev/waveterm/pkg/wshutil"
1720
)
1821

1922
var aiCmd = &cobra.Command{
20-
Use: "ai [-] [message...]",
21-
Short: "Send a message to an AI block",
23+
Use: "ai [options] [files...]",
24+
Short: "Append content to Wave AI sidebar prompt",
25+
Long: `Append content to Wave AI sidebar prompt (does not auto-submit by default)
26+
27+
Arguments:
28+
files... Files to attach (use '-' for stdin)
29+
30+
Examples:
31+
git diff | wsh ai - # Pipe diff to AI, ask question in UI
32+
wsh ai main.go # Attach file, ask question in UI
33+
wsh ai *.go -m "find bugs" # Attach files with message
34+
wsh ai -s - -m "review" < log.txt # Stdin + message, auto-submit
35+
wsh ai -n config.json # New chat with file attached`,
2236
RunE: aiRun,
2337
PreRunE: preRunSetupRpcClient,
2438
DisableFlagsInUseLine: true,
2539
}
2640

27-
var aiFileFlags []string
41+
var aiMessageFlag string
42+
var aiSubmitFlag bool
2843
var aiNewBlockFlag bool
2944

3045
func init() {
3146
rootCmd.AddCommand(aiCmd)
32-
aiCmd.Flags().BoolVarP(&aiNewBlockFlag, "new", "n", false, "create a new AI block")
33-
aiCmd.Flags().StringArrayVarP(&aiFileFlags, "file", "f", nil, "attach file content (use '-' for stdin)")
47+
aiCmd.Flags().StringVarP(&aiMessageFlag, "message", "m", "", "optional message/question to append after files")
48+
aiCmd.Flags().BoolVarP(&aiSubmitFlag, "submit", "s", false, "submit the prompt immediately after appending")
49+
aiCmd.Flags().BoolVarP(&aiNewBlockFlag, "new", "n", false, "create a new AI chat instead of using existing")
3450
}
3551

36-
func encodeFile(builder *strings.Builder, file io.Reader, fileName string) error {
37-
data, err := io.ReadAll(file)
38-
if err != nil {
39-
return fmt.Errorf("error reading file: %w", err)
52+
func detectMimeType(data []byte) string {
53+
mimeType := http.DetectContentType(data)
54+
return strings.Split(mimeType, ";")[0]
55+
}
56+
57+
func getMaxFileSize(mimeType string) (int, string) {
58+
if mimeType == "application/pdf" {
59+
return 5 * 1024 * 1024, "5MB"
4060
}
41-
// Start delimiter with the file name
42-
builder.WriteString(fmt.Sprintf("\n@@@start file %q\n", fileName))
43-
// Read the file content and write it to the builder
44-
builder.Write(data)
45-
// End delimiter with the file name
46-
builder.WriteString(fmt.Sprintf("\n@@@end file %q\n\n", fileName))
47-
return nil
61+
if strings.HasPrefix(mimeType, "image/") {
62+
return 7 * 1024 * 1024, "7MB"
63+
}
64+
return 200 * 1024, "200KB"
4865
}
4966

5067
func aiRun(cmd *cobra.Command, args []string) (rtnErr error) {
5168
defer func() {
5269
sendActivity("ai", rtnErr == nil)
5370
}()
5471

55-
if len(args) == 0 {
72+
if len(args) == 0 && aiMessageFlag == "" {
5673
OutputHelpMessage(cmd)
57-
return fmt.Errorf("no message provided")
74+
return fmt.Errorf("no files or message provided")
5875
}
5976

77+
const maxFileCount = 15
78+
const rpcTimeout = 30000
79+
80+
var allFiles []wshrpc.AIAttachedFile
6081
var stdinUsed bool
61-
var message strings.Builder
6282

63-
// Handle file attachments first
64-
for _, file := range aiFileFlags {
65-
if file == "-" {
83+
if len(args) > maxFileCount {
84+
return fmt.Errorf("too many files (maximum %d files allowed)", maxFileCount)
85+
}
86+
87+
for _, filePath := range args {
88+
var data []byte
89+
var fileName string
90+
var mimeType string
91+
var err error
92+
93+
if filePath == "-" {
6694
if stdinUsed {
6795
return fmt.Errorf("stdin (-) can only be used once")
6896
}
6997
stdinUsed = true
70-
if err := encodeFile(&message, os.Stdin, "<stdin>"); err != nil {
98+
99+
data, err = io.ReadAll(os.Stdin)
100+
if err != nil {
71101
return fmt.Errorf("reading from stdin: %w", err)
72102
}
103+
fileName = "stdin"
104+
mimeType = "text/plain"
73105
} else {
74-
fd, err := os.Open(file)
106+
fileInfo, err := os.Stat(filePath)
75107
if err != nil {
76-
return fmt.Errorf("opening file %s: %w", file, err)
108+
return fmt.Errorf("accessing file %s: %w", filePath, err)
77109
}
78-
defer fd.Close()
79-
if err := encodeFile(&message, fd, file); err != nil {
80-
return fmt.Errorf("reading file %s: %w", file, err)
110+
if fileInfo.IsDir() {
111+
return fmt.Errorf("%s is a directory, not a file", filePath)
81112
}
82-
}
83-
}
84113

85-
// Default to "waveai" block
86-
isDefaultBlock := blockArg == ""
87-
if isDefaultBlock {
88-
blockArg = "view@waveai"
89-
}
90-
var fullORef *waveobj.ORef
91-
var err error
92-
if !aiNewBlockFlag {
93-
fullORef, err = resolveSimpleId(blockArg)
94-
}
95-
if (err != nil && isDefaultBlock) || aiNewBlockFlag {
96-
// Create new AI block if default block doesn't exist
97-
data := &wshrpc.CommandCreateBlockData{
98-
BlockDef: &waveobj.BlockDef{
99-
Meta: map[string]interface{}{
100-
waveobj.MetaKey_View: "waveai",
101-
},
102-
},
103-
Focused: true,
114+
data, err = os.ReadFile(filePath)
115+
if err != nil {
116+
return fmt.Errorf("reading file %s: %w", filePath, err)
117+
}
118+
fileName = filepath.Base(filePath)
119+
mimeType = detectMimeType(data)
104120
}
105121

106-
newORef, err := wshclient.CreateBlockCommand(RpcClient, *data, &wshrpc.RpcOpts{Timeout: 2000})
107-
if err != nil {
108-
return fmt.Errorf("creating AI block: %w", err)
109-
}
110-
fullORef = &newORef
111-
// Wait for the block's route to be available
112-
gotRoute, err := wshclient.WaitForRouteCommand(RpcClient, wshrpc.CommandWaitForRouteData{
113-
RouteId: wshutil.MakeFeBlockRouteId(fullORef.OID),
114-
WaitMs: 4000,
115-
}, &wshrpc.RpcOpts{Timeout: 5000})
116-
if err != nil {
117-
return fmt.Errorf("waiting for AI block: %w", err)
122+
isPDF := mimeType == "application/pdf"
123+
isImage := strings.HasPrefix(mimeType, "image/")
124+
125+
if !isPDF && !isImage {
126+
mimeType = "text/plain"
127+
if utilfn.ContainsBinaryData(data) {
128+
return fmt.Errorf("file %s contains binary data and cannot be uploaded as text", fileName)
129+
}
118130
}
119-
if !gotRoute {
120-
return fmt.Errorf("AI block route could not be established")
131+
132+
maxSize, sizeStr := getMaxFileSize(mimeType)
133+
if len(data) > maxSize {
134+
return fmt.Errorf("file %s exceeds maximum size of %s for %s files", fileName, sizeStr, mimeType)
121135
}
122-
} else if err != nil {
123-
return fmt.Errorf("resolving block: %w", err)
136+
137+
allFiles = append(allFiles, wshrpc.AIAttachedFile{
138+
Name: fileName,
139+
Type: mimeType,
140+
Size: len(data),
141+
Data64: base64.StdEncoding.EncodeToString(data),
142+
})
124143
}
125144

126-
// Create the route for this block
127-
route := wshutil.MakeFeBlockRouteId(fullORef.OID)
145+
tabId := os.Getenv("WAVETERM_TABID")
146+
if tabId == "" {
147+
return fmt.Errorf("WAVETERM_TABID environment variable not set")
148+
}
149+
150+
route := wshutil.MakeTabRouteId(tabId)
128151

129-
// Then handle main message
130-
if args[0] == "-" {
131-
if stdinUsed {
132-
return fmt.Errorf("stdin (-) can only be used once")
152+
if aiNewBlockFlag {
153+
newChatData := wshrpc.CommandWaveAIAddContextData{
154+
NewChat: true,
133155
}
134-
data, err := io.ReadAll(os.Stdin)
156+
err := wshclient.WaveAIAddContextCommand(RpcClient, newChatData, &wshrpc.RpcOpts{
157+
Route: route,
158+
Timeout: rpcTimeout,
159+
})
135160
if err != nil {
136-
return fmt.Errorf("reading from stdin: %w", err)
137-
}
138-
message.Write(data)
139-
140-
// Also include any remaining arguments (excluding the "-" itself)
141-
if len(args) > 1 {
142-
if message.Len() > 0 {
143-
message.WriteString(" ")
144-
}
145-
message.WriteString(strings.Join(args[1:], " "))
161+
return fmt.Errorf("creating new chat: %w", err)
146162
}
147-
} else {
148-
message.WriteString(strings.Join(args, " "))
149163
}
150164

151-
if message.Len() == 0 {
152-
return fmt.Errorf("message is empty")
153-
}
154-
if message.Len() > 50*1024 {
155-
return fmt.Errorf("current max message size is 50k")
165+
for _, file := range allFiles {
166+
contextData := wshrpc.CommandWaveAIAddContextData{
167+
Files: []wshrpc.AIAttachedFile{file},
168+
}
169+
err := wshclient.WaveAIAddContextCommand(RpcClient, contextData, &wshrpc.RpcOpts{
170+
Route: route,
171+
Timeout: rpcTimeout,
172+
})
173+
if err != nil {
174+
return fmt.Errorf("adding file %s: %w", file.Name, err)
175+
}
156176
}
157177

158-
messageData := wshrpc.AiMessageData{
159-
Message: message.String(),
160-
}
161-
err = wshclient.AiSendMessageCommand(RpcClient, messageData, &wshrpc.RpcOpts{
162-
Route: route,
163-
Timeout: 2000,
164-
})
165-
if err != nil {
166-
return fmt.Errorf("sending message: %w", err)
178+
if aiMessageFlag != "" || aiSubmitFlag {
179+
finalContextData := wshrpc.CommandWaveAIAddContextData{
180+
Text: aiMessageFlag,
181+
Submit: aiSubmitFlag,
182+
}
183+
err := wshclient.WaveAIAddContextCommand(RpcClient, finalContextData, &wshrpc.RpcOpts{
184+
Route: route,
185+
Timeout: rpcTimeout,
186+
})
187+
if err != nil {
188+
return fmt.Errorf("adding context: %w", err)
189+
}
167190
}
168191

169192
return nil

‎docs/docs/wsh-reference.mdx‎

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -110,25 +110,45 @@ wsh getmeta -b [other-tab-id] "bg:*" --clear-prefix | wsh setmeta -b tab --json
110110
111111
## ai
112112
113-
Send messages to new or existing AI blocks directly from the CLI. `-f` passes a file. note that there is a maximum size of 10k for messages and files, so use a tail/grep to cut down file sizes before passing. The `-f` option works great for small files though like shell scripts or `.zshrc` etc. You can use "-" to read input from stdin.
113+
Append content to the Wave AI sidebar. Files are attached as proper file attachments (supporting images, PDFs, and text), not encoded as text. By default, content is added to the sidebar without auto-submitting, allowing you to review and add more context before sending to the AI.
114114
115-
By default the messages get sent to the first AI block (by blocknum). If no AI block exists, then a new one will be created. Use `-n` to force creation of a new AI block. Use `-b` to target a specific AI block.
115+
You can attach multiple files at once (up to 15 files). Use `-m` to add a message along with files, `-s` to auto-submit immediately, and `-n` to start a new chat conversation. Use "-" to read from stdin.
116116
117117
```sh
118-
wsh ai "how do i write an ls command that sorts files in reverse size order"
119-
wsh ai -f <(tail -n 20 "my.log") -- "any idea what these error messages mean"
120-
wsh ai -f README.md "help me update this readme file"
118+
# Pipe command output to AI (ask question in UI)
119+
git diff | wsh ai -
120+
docker logs mycontainer | wsh ai -
121121
122-
# creates a new AI block
123-
wsh ai -n "tell me a story"
122+
# Attach files without auto-submit (review in UI first)
123+
wsh ai main.go utils.go
124+
wsh ai screenshot.png logs.txt
124125
125-
# targets block number 5
126-
wsh ai -b 5 "tell me more"
126+
# Attach files with message
127+
wsh ai app.py -m "find potential bugs"
128+
wsh ai *.log -m "analyze these error logs"
127129
128-
# read from stdin and also supply a message
129-
tail -n 50 mylog.log | wsh ai - "can you tell me what this error means?"
130+
# Auto-submit immediately
131+
wsh ai config.json -s -m "explain this configuration"
132+
tail -n 50 app.log | wsh ai -s - -m "what's causing these errors?"
133+
134+
# Start new chat and attach files
135+
wsh ai -n report.pdf data.csv -m "summarize these reports"
136+
137+
# Attach different file types (images, PDFs, code)
138+
wsh ai architecture.png api-spec.pdf server.go -m "review the system design"
130139
```
131140
141+
**File Size Limits:**
142+
- Text files: 200KB maximum
143+
- PDF files: 5MB maximum
144+
- Image files: 7MB maximum (accounts for base64 encoding overhead)
145+
- Maximum 15 files per command
146+
147+
**Flags:**
148+
- `-m, --message <text>` - Add message text along with files
149+
- `-s, --submit` - Auto-submit immediately (default waits for user)
150+
- `-n, --new` - Clear current chat and start fresh conversation
151+
132152
---
133153
134154
## editconfig

‎docs/docs/wsh.mdx‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -116,17 +116,40 @@ wsh setvar -b tab SHARED_ENV=staging
116116

117117
### AI-Assisted Development
118118

119+
The `wsh ai` command appends content to the Wave AI sidebar. By default, files are attached without auto-submitting, allowing you to review and add more context before sending.
120+
119121
```bash
120-
# Get AI help with code (uses "-" to read from stdin)
121-
git diff | wsh ai - "review these changes"
122+
# Pipe output to AI sidebar (ask question in UI)
123+
git diff | wsh ai -
124+
125+
# Attach files with a message
126+
wsh ai main.go utils.go -m "find bugs in these files"
127+
128+
# Auto-submit with message
129+
wsh ai config.json -s -m "explain this config"
122130

123-
# Get help with a file
124-
wsh ai -f .zshrc "help me add ~/bin to my path"
131+
# Start new chat with attached files
132+
wsh ai -n *.log -m "analyze these logs"
125133

126-
# Debug issues (uses "-" to read from stdin)
127-
dmesg | wsh ai - "help me understand these errors"
134+
# Attach multiple file types (images, PDFs, code)
135+
wsh ai screenshot.png report.pdf app.py -m "review these"
136+
137+
# Debug with stdin and auto-submit
138+
dmesg | wsh ai -s - -m "help me understand these errors"
128139
```
129140

141+
**Flags:**
142+
- `-` - Read from stdin instead of a file
143+
- `-m, --message` - Add message text along with files
144+
- `-s, --submit` - Auto-submit immediately (default is to wait for user)
145+
- `-n, --new` - Clear chat and start fresh conversation
146+
147+
**File Limits:**
148+
- Text files: 200KB max
149+
- PDFs: 5MB max
150+
- Images: 7MB max
151+
- Maximum 15 files per command
152+
130153
## Tips & Features
131154

132155
1. **Working with Blocks**

0 commit comments

Comments
 (0)