feat(tool): update web_search_aichat.sh (#166)

1 files changed, 11 insertions, 4 deletions

diff --git a/tools/web_search_aichat.sh b/tools/web_search_aichat.sh
index a20bd35..8587e71 100755
--- a/tools/web_search_aichat.sh
+++ b/tools/web_search_aichat.sh
@@ -11,16 +11,23 @@ set -e
 # @env WEB_SEARCH_MODEL! The model for web-searching.
 #
 # supported aichat models:
+#   - gemini:gemini-2.0-*
 #   - vertexai:gemini-*
-#   - perplexity:*-online
+#   - perplexity:*
 #   - ernie:*
-#   - lingyiwanwu:yi-large-rag
 # @env LLM_OUTPUT=/dev/stdout The output path
 
 main() {
     client="${WEB_SEARCH_MODEL%%:*}"
-    if [[ "$client" == "vertexai" ]]; then
-        export AICHAT_PATCH_VERTEXAI_CHAT_COMPLETIONS='{"gemini-.*":{"body":{"tools":[{"googleSearchRetrieval":{}}]}}}'
+    if [[ "$client" == "gemini" ]]; then
+        export AICHAT_PATCH_GEMINI_CHAT_COMPLETIONS='{".*":{"body":{"tools":[{"google_search":{}}]}}}'
+    elif [[ "$client" == "vertexai" ]]; then
+        export AICHAT_PATCH_VERTEXAI_CHAT_COMPLETIONS='{
+    "gemini-1.5-.*":{"body":{"tools":[{"googleSearchRetrieval":{}}]}},
+    "gemini-2.0-.*":{"body":{"tools":[{"google_search":{}}]}}
+}'
+    elif [[ "$client" == "ernie" ]]; then
+        export AICHAT_PATCH_ERNIE_CHAT_COMPLETIONS='{".*":{"body":{"web_search":{"enable":true}}}}'
     fi
     aichat -m "$WEB_SEARCH_MODEL" "$argc_query" >> "$LLM_OUTPUT"
 }