255 lines
7.1 KiB
Bash
Executable File
255 lines
7.1 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -e
|
|
|
|
# @describe Perform a web search to get up-to-date information or additional context.
|
|
# Use this when you need current information or feel a search could provide a better answer.
|
|
|
|
# @option --query! The search query.
|
|
|
|
# @meta require-tools loki
|
|
|
|
# @env WEB_SEARCH_MODEL=gemini:gemini-2.5-flash The model for web-searching.
|
|
#
|
|
# supported loki models:
|
|
# - gemini:gemini-2.0-*
|
|
# - vertexai:gemini-*
|
|
# - perplexity:*
|
|
# - ernie:*
|
|
# @env LLM_OUTPUT=/dev/stdout The output path
|
|
|
|
# shellcheck disable=SC2154
|
|
main() {
|
|
client="${WEB_SEARCH_MODEL%%:*}"
|
|
|
|
if [[ "$client" == "gemini" ]]; then
|
|
export LOKI_PATCH_GEMINI_CHAT_COMPLETIONS='{".*":{"body":{"tools":[{"google_search":{}}]}}}'
|
|
elif [[ "$client" == "vertexai" ]]; then
|
|
export LOKI_PATCH_VERTEXAI_CHAT_COMPLETIONS='{
|
|
"gemini-1.5-.*":{"body":{"tools":[{"googleSearchRetrieval":{}}]}},
|
|
"gemini-2.0-.*":{"body":{"tools":[{"google_search":{}}]}}
|
|
}'
|
|
elif [[ "$client" == "ernie" ]]; then
|
|
export LOKI_PATCH_ERNIE_CHAT_COMPLETIONS='{".*":{"body":{"web_search":{"enable":true}}}}'
|
|
fi
|
|
|
|
loki -m "$WEB_SEARCH_MODEL" "$argc_query" >> "$LLM_OUTPUT"
|
|
}
|
|
|
|
|
|
# ARGC-BUILD {
|
|
# This block was generated by argc (https://github.com/sigoden/argc).
|
|
# Modifying it manually is not recommended
|
|
|
|
_argc_run() {
|
|
if [[ "${1:-}" == "___internal___" ]]; then
|
|
_argc_die "error: unsupported ___internal___ command"
|
|
fi
|
|
if [[ "${OS:-}" == "Windows_NT" ]] && [[ -n "${MSYSTEM:-}" ]]; then
|
|
set -o igncr
|
|
fi
|
|
argc__args=("$(basename "$0" .sh)" "$@")
|
|
argc__positionals=()
|
|
_argc_index=1
|
|
_argc_len="${#argc__args[@]}"
|
|
_argc_tools=()
|
|
_argc_parse
|
|
_argc_require_tools "${_argc_tools[@]}"
|
|
if [ -n "${argc__fn:-}" ]; then
|
|
$argc__fn "${argc__positionals[@]}"
|
|
fi
|
|
}
|
|
|
|
_argc_usage() {
|
|
cat <<-'EOF'
|
|
Perform a web search to get up-to-date information or additional context.
|
|
Use this when you need current information or feel a search could provide a better answer.
|
|
|
|
USAGE: --query <QUERY>
|
|
|
|
OPTIONS:
|
|
--query <QUERY> The query to search for.
|
|
-h, --help Print help
|
|
-V, --version Print version
|
|
|
|
ENVIRONMENTS:
|
|
WEB_SEARCH_MODEL
|
|
The model for web-searching.
|
|
|
|
supported loki models:
|
|
- gemini:gemini-2.0-*
|
|
- vertexai:gemini-*
|
|
- perplexity:*
|
|
- ernie:*
|
|
[default: gemini:gemini-2.5-flash]
|
|
|
|
LLM_OUTPUT
|
|
The output path [default: /dev/stdout]
|
|
EOF
|
|
exit
|
|
}
|
|
|
|
_argc_version() {
|
|
echo 0.0.0
|
|
exit
|
|
}
|
|
|
|
_argc_parse() {
|
|
local _argc_key _argc_action
|
|
local _argc_subcmds=""
|
|
while [[ $_argc_index -lt $_argc_len ]]; do
|
|
_argc_item="${argc__args[_argc_index]}"
|
|
_argc_key="${_argc_item%%=*}"
|
|
case "$_argc_key" in
|
|
--help | -help | -h)
|
|
_argc_usage
|
|
;;
|
|
--version | -version | -V)
|
|
_argc_version
|
|
;;
|
|
--)
|
|
_argc_dash="${#argc__positionals[@]}"
|
|
argc__positionals+=("${argc__args[@]:$((_argc_index + 1))}")
|
|
_argc_index=$_argc_len
|
|
break
|
|
;;
|
|
--query)
|
|
_argc_take_args "--query <QUERY>" 1 1 "-" ""
|
|
_argc_index=$((_argc_index + _argc_take_args_len + 1))
|
|
if [[ -z "${argc_query:-}" ]]; then
|
|
argc_query="${_argc_take_args_values[0]:-}"
|
|
else
|
|
_argc_die "error: the argument \`--query\` cannot be used multiple times"
|
|
fi
|
|
;;
|
|
*)
|
|
if _argc_maybe_flag_option "-" "$_argc_item"; then
|
|
_argc_die "error: unexpected argument \`$_argc_key\` found"
|
|
fi
|
|
argc__positionals+=("$_argc_item")
|
|
_argc_index=$((_argc_index + 1))
|
|
;;
|
|
esac
|
|
done
|
|
_argc_require_params "error: the following required arguments were not provided:" \
|
|
'argc_query:--query <QUERY>'
|
|
_argc_tools=(loki)
|
|
if [[ -n "${_argc_action:-}" ]]; then
|
|
$_argc_action
|
|
else
|
|
argc__fn=main
|
|
if [[ "${argc__positionals[0]:-}" == "help" ]] && [[ "${#argc__positionals[@]}" -eq 1 ]]; then
|
|
_argc_usage
|
|
fi
|
|
if [[ -z "${WEB_SEARCH_MODEL:-}" ]]; then
|
|
export WEB_SEARCH_MODEL=gemini:gemini-2.5-flash
|
|
fi
|
|
if [[ -z "${LLM_OUTPUT:-}" ]]; then
|
|
export LLM_OUTPUT=/dev/stdout
|
|
fi
|
|
fi
|
|
}
|
|
|
|
_argc_take_args() {
|
|
_argc_take_args_values=()
|
|
_argc_take_args_len=0
|
|
local param="$1" min="$2" max="$3" signs="$4" delimiter="$5"
|
|
if [[ "$min" -eq 0 ]] && [[ "$max" -eq 0 ]]; then
|
|
return
|
|
fi
|
|
local _argc_take_index=$((_argc_index + 1)) _argc_take_value
|
|
if [[ "$_argc_item" == *=* ]]; then
|
|
_argc_take_args_values=("${_argc_item##*=}")
|
|
else
|
|
while [[ $_argc_take_index -lt $_argc_len ]]; do
|
|
_argc_take_value="${argc__args[_argc_take_index]}"
|
|
if _argc_maybe_flag_option "$signs" "$_argc_take_value"; then
|
|
if [[ "${#_argc_take_value}" -gt 1 ]]; then
|
|
break
|
|
fi
|
|
fi
|
|
_argc_take_args_values+=("$_argc_take_value")
|
|
_argc_take_args_len=$((_argc_take_args_len + 1))
|
|
if [[ "$_argc_take_args_len" -ge "$max" ]]; then
|
|
break
|
|
fi
|
|
_argc_take_index=$((_argc_take_index + 1))
|
|
done
|
|
fi
|
|
if [[ "${#_argc_take_args_values[@]}" -lt "$min" ]]; then
|
|
_argc_die "error: incorrect number of values for \`$param\`"
|
|
fi
|
|
if [[ -n "$delimiter" ]] && [[ "${#_argc_take_args_values[@]}" -gt 0 ]]; then
|
|
local item values arr=()
|
|
for item in "${_argc_take_args_values[@]}"; do
|
|
IFS="$delimiter" read -r -a values <<<"$item"
|
|
arr+=("${values[@]}")
|
|
done
|
|
_argc_take_args_values=("${arr[@]}")
|
|
fi
|
|
}
|
|
|
|
_argc_require_params() {
|
|
local message="$1" missed_envs="" item name render_name
|
|
for item in "${@:2}"; do
|
|
name="${item%%:*}"
|
|
render_name="${item##*:}"
|
|
if [[ -z "${!name:-}" ]]; then
|
|
missed_envs="$missed_envs"$'\n'" $render_name"
|
|
fi
|
|
done
|
|
if [[ -n "${missed_envs}" ]]; then
|
|
_argc_die "$message$missed_envs"
|
|
fi
|
|
}
|
|
|
|
_argc_maybe_flag_option() {
|
|
local signs="$1" arg="$2"
|
|
if [[ -z "$signs" ]]; then
|
|
return 1
|
|
fi
|
|
local cond=false
|
|
if [[ "$signs" == *"+"* ]]; then
|
|
if [[ "$arg" =~ ^\+[^+].* ]]; then
|
|
cond=true
|
|
fi
|
|
elif [[ "$arg" == -* ]]; then
|
|
if (( ${#arg} < 3 )) || [[ ! "$arg" =~ ^---.* ]]; then
|
|
cond=true
|
|
fi
|
|
fi
|
|
if [[ "$cond" == "false" ]]; then
|
|
return 1
|
|
fi
|
|
local value="${arg%%=*}"
|
|
if [[ "$value" =~ [[:space:]] ]]; then
|
|
return 1
|
|
fi
|
|
return 0
|
|
}
|
|
|
|
_argc_require_tools() {
|
|
local tool missing_tools=()
|
|
for tool in "$@"; do
|
|
if ! command -v "$tool" >/dev/null 2>&1; then
|
|
missing_tools+=("$tool")
|
|
fi
|
|
done
|
|
if [[ "${#missing_tools[@]}" -gt 0 ]]; then
|
|
echo "error: missing tools: ${missing_tools[*]}" >&2
|
|
exit 1
|
|
fi
|
|
}
|
|
|
|
_argc_die() {
|
|
if [[ $# -eq 0 ]]; then
|
|
cat
|
|
else
|
|
echo "$*" >&2
|
|
fi
|
|
exit 1
|
|
}
|
|
|
|
_argc_run "$@"
|
|
|
|
# ARGC-BUILD }
|