feat(tool): add web_search_aichat.sh (#120)

2024-11-02 20:42:42 +08:00
parent 68b1f57a4b
commit 3573a43978
3 changed files with 32 additions and 78 deletions
@@ -0,0 +1,32 @@
+#!/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 query to search for.
+
+# @meta require-tools aichat
+# @env WEB_SEARCH_MODEL! The model for web-searching.
+#
+# supported aichat models:
+#   - cohere:*
+#   - vertexai:gemini-*
+#   - perplexity:*-online
+#   - ernie:*
+#   - lingyiwanwu:yi-large-rag
+
+main() {
+    client="${WEB_SEARCH_MODEL%%:*}"
+    case "$client" in
+    cohere)
+        export AICHAT_PATCH_COHERE_CHAT_COMPLETIONS='{".*":{"body":{"connectors":[{"id":"web-search"}]}}}'
+        ;;
+    vertexai)
+        export AICHAT_PATCH_VERTEXAI_CHAT_COMPLETIONS='{"gemini-.*":{"body":{"tools":[{"googleSearchRetrieval":{}}]}}}'
+        ;;
+    esac
+    aichat -m "$WEB_SEARCH_MODEL" "$argc_query" >> "$LLM_OUTPUT"
+}
+
+eval "$(argc --argc-eval "$0" "$@")"
@@ -1,27 +0,0 @@
-#!/usr/bin/env bash
-set -e
-
-# @describe Perform a web search using Cohere API 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.
-
-# @env COHERE_API_KEY! The api key
-# @env COHERE_WEB_SEARCH_MODEL=command-r The LLM model for web search
-# @option --query! The query to search for.
-
-main() {
-    curl -fsS -X POST https://api.cohere.com/v1/chat \
-     -H "authorization: Bearer $COHERE_API_KEY" \
-     -H "accept: application/json" \
-     -H "content-type: application/json" \
-     --data '
-{
-    "model": "'"$COHERE_WEB_SEARCH_MODEL"'",
-    "message": "'"$argc_query"'",
-    "connectors": [{"id": "web-search"}]
-}
-'  | \
-        jq -r '.text' \
-        >> "$LLM_OUTPUT"
-}
-
-eval "$(argc --argc-eval "$0" "$@")"
@@ -1,51 +0,0 @@
-#!/usr/bin/env bash
-set -e
-
-# @describe Perform a web search using VertexAI Gemini API 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.
-
-# @env VERTEXAI_PROJECT_ID! The project id
-# @env VERTEXAI_LOCATION! The location
-# @env VERTEXAI_WEB_SEARCH_MODEL=gemini-1.5-pro-001 The LLM model for web search
-# @option --query! The query to search for.
-# @meta require-tools gcloud
-
-main() {
-    curl -fsSL https://$VERTEXAI_LOCATION-aiplatform.googleapis.com/v1beta1/projects/$VERTEXAI_PROJECT_ID/locations/$VERTEXAI_LOCATION/publishers/google/models/$VERTEXAI_WEB_SEARCH_MODEL:generateContent \
-        -X POST \
-        -H "Authorization: Bearer $(gcloud auth application-default print-access-token)" \
-        -H "Content-Type: application/json" \
-        -d '
-{
-    "contents": [{
-        "role": "user",
-        "parts": [{
-            "text": "'"$argc_query"'"
-        }]
-    }],
-    "safetySettings": [
-        {
-            "category": "HARM_CATEGORY_HARASSMENT",
-            "threshold": "BLOCK_ONLY_HIGH"
-        },
-        {
-            "category": "HARM_CATEGORY_HATE_SPEECH",
-            "threshold": "BLOCK_ONLY_HIGH"
-        },
-        {
-            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-            "threshold": "BLOCK_ONLY_HIGH"
-        },
-        {
-            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-            "threshold": "BLOCK_ONLY_HIGH"
-        }
-    ],
-    "tools": [{
-        "googleSearchRetrieval": {}
-    }]
-  }' | \
-    jq -r '.candidates[0].content.parts[0].text' >> "$LLM_OUTPUT"
-}
-
-eval "$(argc --argc-eval "$0" "$@")"