-
Notifications
You must be signed in to change notification settings - Fork 2
/
gpt
executable file
·354 lines (282 loc) · 8.66 KB
/
gpt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
#!/usr/bin/env bash
set -eu -o pipefail
# ------------------------------------------------------------------------------
# Global variables
# ------------------------------------------------------------------------------
OPENAI_MODEL="${OPENAI_MODEL:-gpt-3.5-turbo-16k}"
BASHGPT_DEBUG="${BASHGPT_DEBUG:-false}"
PROMPTS=()
USER_PROMPTS=()
SYSTEM_PROMPTS=()
OPTIONS=()
READ_STDIN=false
SLURP_STDIN=false
PAYLOAD_FILE=$(mktemp)
CACHE_DIR="${BASHGPT_CACHE_DIR:-/$HOME/.bashgpt_cache}"
CLEAR_CACHE=false
CLEAR_CACHE_ALL=false
NO_CACHE=false
# Ensure that the temp file holding the payload is deleted on exit
trap 'cleanup' EXIT ERR INT
# ------------------------------------------------------------------------------
# Functions
# ------------------------------------------------------------------------------
usage() {
local exit_code="${1:-0}"
cat << EOL
usage: gpt [OPTIONS]
Options:
--help | -h show help information
--user | -u add a user prompt to the list; can be used multiple times
--system | -s add a system prompt to the list; can be used multiple times
--prompt | -p add a prompt to the list; can be used multiple times (deprecated; use -u instead)
--option | -o set an openai chat completions API option; can be used multiple times
--read | -r read prompts from STDIN
--slurp with -r, read STDIN until EOF and perform a single GPT request with the input;
if unset, each line of STDIN will result in a new GPT request
--clear | -c clear the cache for the current prompt
--clear-all clear the cache for all prompts
--no-cache do not use the cache
--check check for required tools and environment variables
Environmental variables:
OPENAI_API_KEY your OpenAI API key (required; create one at https://platform.openai.com/account/api-keys)
OPENAI_MODEL the model to use; defaults to gpt-3.5-turbo-16k
BASHGPT_CACHE_DIR the directory to store cached results; defaults to \$HOME/.bashgpt_cache
Examples:
# Generate a ridiculous script in bash, format it using gum
gpt -u 'Write a skewheap in bash' | gum format --theme dark
EOL
exit "$exit_code"
}
cleanup() {
rm -f "$PAYLOAD_FILE"
}
validate_env() {
local tools=("curl" "jq" "awk" "sort" "sha256sum" "tee")
local cmd
for cmd in "${tools[@]}"; do
if ! command -v "$cmd" &> /dev/null; then
echo >&2 "Error: $cmd is not installed."
exit 1
fi
done
if [ -z "${OPENAI_API_KEY:-}" ]; then
echo >&2 "Error: OPENAI_API_KEY is not set in the environment."
exit 1
fi
return 0
}
make_cache_dir() {
mkdir -p "$CACHE_DIR" 2> /dev/null || {
echo >&2 "Failed to create cache directory '$CACHE_DIR'"
exit 1
}
return 0
}
clear_cache() {
rm -f "$CACHE_DIR"/* || {
echo >&2 "Error clearing cache directory '$CACHE_DIR'"
exit 1
}
return 0
}
clear_cache_key() {
local cache_key="$1"
local cache_file="$CACHE_DIR/$cache_key"
rm -f "$cache_file" || {
echo >&2 "Failed to clear cache for current prompt"
exit 1
}
return 0
}
cache_key() {
local sorted_options
local sorted_system_prompts
local sorted_user_prompts
local sorted_prompts
IFS=$'\n' mapfile -t sorted_options < <(sort <<< "${OPTIONS[*]}")
IFS=$'\n' mapfile -t sorted_system_prompts < <(sort <<< "${SYSTEM_PROMPTS[*]}")
IFS=$'\n' mapfile -t sorted_user_prompts < <(sort <<< "${USER_PROMPTS[*]}")
IFS=$'\n' mapfile -t sorted_prompts < <(sort <<< "${PROMPTS[*]}")
echo -n "${sorted_options[*]}${sorted_system_prompts[*]}${sorted_user_prompts[*]}${sorted_prompts[*]}" \
| sha256sum | awk '{print $1}'
}
inspect() {
local line
while IFS= read -r line; do
if [ "$BASHGPT_DEBUG" = "true" ]; then
gum >&2 join "$(gum style --foreground="#00FF00" 'inspect> ')" "$line"
fi
echo "${line}"
done
}
build_payload() {
local payload
local new_message
# Initialize payload with model and stream fields
payload=$(jq -n --arg model "$OPENAI_MODEL" '{"model": $model, "stream": true, "messages": []}')
for option in "${OPTIONS[@]}"; do
key="${option%%:*}"
key="${key#"${key%%[![:space:]]*}"}"
key="${key%"${key##*[![:space:]]}"}"
val="${option#*:}"
val="${val#"${val%%[![:space:]]*}"}"
val="${val%"${val##*[![:space:]]}"}"
payload=$(jq --arg k "$key" --arg v "$val" '. + {($k): $v}' <<< "$payload")
done
# Loop through system_prompts to add more messages
for prompt in "${SYSTEM_PROMPTS[@]}"; do
new_message=$(jq -n --arg content "$prompt" '{"role": "system", "content": $content}')
payload=$(jq --argjson new_msg "$new_message" '.messages += [$new_msg]' <<< "$payload")
done
# Loop through user_prompts to add more messages
for prompt in "${USER_PROMPTS[@]}"; do
new_message=$(jq -n --arg content "$prompt" '{"role": "user", "content": $content}')
payload=$(jq --argjson new_msg "$new_message" '.messages += [$new_msg]' <<< "$payload")
done
# Loop through all_prompts to add more messages
for prompt in "${PROMPTS[@]}"; do
new_message=$(jq -n --arg content "$prompt" '{"role": "user", "content": $content}')
payload=$(jq --argjson new_msg "$new_message" '.messages += [$new_msg]' <<< "$payload")
done
echo "$payload" > "$PAYLOAD_FILE"
}
get_completion() {
local exit_code
build_payload
curl -s \
-N "https://api.openai.com/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $OPENAI_API_KEY" \
--data-binary "@$PAYLOAD_FILE" \
| inspect \
| while IFS= read -r line; do
if [[ $line == *"[DONE]"* ]]; then
exit 0
fi
processed_line=$(echo "$line" | awk '/data: / { sub(/^data: /, ""); print }')
if [[ -n $processed_line ]]; then
if echo "$processed_line" | jq -e 'has("error")' > /dev/null 2>&1; then
error_message=$(echo "$processed_line" | jq -r '.error.message')
echo -e "\e[31m${error_message}\e[0m" # Print in red
exit 1
else
echo "$processed_line" | jq -j 'if .choices[0].delta.content == null then "\n" else .choices[0].delta.content end'
fi
fi
done
exit_code=$?
if [ $exit_code -ne 0 ]; then
echo >&2 "Error: curl request failed with exit code $exit_code."
exit 1
fi
}
single_completion() {
# Build cache key and cache file path
cache_key=$(cache_key)
cache_file="$CACHE_DIR/$cache_key"
# Clear the cache for this request if --clear was set
if [ "$CLEAR_CACHE" = "true" ]; then
clear_cache_key "$cache_key"
fi
# Get our completion
if [ "$NO_CACHE" = "true" ]; then
get_completion
elif [ -e "$cache_file" ] && [ -s "$cache_file" ]; then
cat "$cache_file"
else
# Use stdbuf to ensure that the output is not buffered, since get_completion
# is not printing by line.
get_completion | stdbuf -o0 tee -a "$cache_file"
fi
}
slurpy_completion() {
mapfile -t lines
USER_PROMPTS=("${lines[@]}" "${USER_PROMPTS[@]}")
single_completion
}
readliney_completion() {
original=("${USER_PROMPTS[@]}")
while IFS= read -r line; do
USER_PROMPTS+=("$line")
single_completion
USER_PROMPTS=("${original[@]}")
done
}
# ------------------------------------------------------------------------------
# Process arguments
# ------------------------------------------------------------------------------
while (("$#")); do
case "$1" in
--help | -h)
usage 0
;;
--opt | -o)
OPTIONS+=("$2")
shift 2
;;
--user | -u)
USER_PROMPTS+=("$2")
shift 2
;;
--system | -s)
SYSTEM_PROMPTS+=("$2")
shift 2
;;
--prompt | -p)
PROMPTS+=("$2")
shift 2
;;
--read | -r)
READ_STDIN=true
shift
;;
--slurp)
SLURP_STDIN=true
shift
;;
--no-cache)
NO_CACHE=true
shift
;;
--clear | -c)
CLEAR_CACHE=true
shift
;;
--clear-all)
CLEAR_CACHE_ALL=true
shift
;;
--check)
validate_env
exit 0
;;
*)
echo >&2 "Error: unknown argument '$1'"
usage 1
;;
esac
done
# ------------------------------------------------------------------------------
# Main logic
# ------------------------------------------------------------------------------
# Validate shell environment and tools
validate_env
# Create the cache directory if it doesn't exist
make_cache_dir
# Clear all cached results if --clear-all was set
if [ "$CLEAR_CACHE_ALL" = "true" ]; then
clear_cache
exit 0
fi
# Check if there's input from STDIN
if [ "$READ_STDIN" = "true" ] && [[ ! -t 0 ]]; then
if [ "$SLURP_STDIN" = "true" ]; then
slurpy_completion
else
readliney_completion
fi
else
single_completion
fi
echo