Commit d4cfff1

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-02-11 05:14:54
chore(gpt): update search engine fallback
1 parent 041eccf
Changed files (2)
src/llm/tool_call.py
@@ -12,11 +12,38 @@ from networking import hx_req
 
 async def get_online_search_result(query: str) -> list[dict]:
     if GPT.PRIMARY_SEARCH_ENGINE == "google":
-        return await google_search_result(query)
+        results = await google_search(query)
+        if not results:
+            return await glm_search(query)
+    elif GPT.PRIMARY_SEARCH_ENGINE == "glm":
+        results = await glm_search(query)
+        if not results:
+            return await google_search(query)
+    return []
+
+
+async def google_search(query: str) -> list[dict]:
+    try:
+        url = f"https://www.googleapis.com/customsearch/v1?key={TOKEN.GOOGLE_SEARCH_API_KEY}&cx={TOKEN.GOOGLE_SEARCH_CX}&q={query}"
+        response = await hx_req(url, proxy=PROXY.GOOGLE_SEARCH, check_keys=["items"])
+        results = glom(response, "items", default=[]) or []
+        for item in results:
+            keys = copy.copy(item).keys()
+            for key in keys:
+                if key not in ["title", "link", "snippet", "mime"]:
+                    item.pop(key, None)
+        if results:
+            return results[: int(GPT.SEARCH_NUM_RESULTS)]
+    except Exception as e:
+        logger.error(e)
+    return []
+
+
+async def glm_search(query: str) -> list[dict]:
     try:
         client = AsyncOpenAI(
-            api_key=GPT.SEARCH_API_KEY,
-            base_url=GPT.SEARCH_BASE_URL,
+            api_key=GPT.GLM_API_KEY,
+            base_url=GPT.GLM_BASE_URL,
             http_client=DefaultAsyncHttpxClient(proxy=PROXY.GPT),
         )
         tools = [{"type": "web_search", "web_search": {"enable": True, "search_query": query, "search_result": True}}]
@@ -37,21 +64,4 @@ async def get_online_search_result(query: str) -> list[dict]:
             return results[: int(GPT.SEARCH_NUM_RESULTS)]
     except Exception as e:
         logger.error(e)
-    return await google_search_result(query)
-
-
-async def google_search_result(query: str) -> list[dict]:
-    try:
-        url = f"https://www.googleapis.com/customsearch/v1?key={TOKEN.GOOGLE_SEARCH_API_KEY}&cx={TOKEN.GOOGLE_SEARCH_CX}&q={query}"
-        response = await hx_req(url, proxy=PROXY.GOOGLE_SEARCH, check_keys=["items"])
-        results = glom(response, "items", default=[]) or []
-        for item in results:
-            keys = copy.copy(item).keys()
-            for key in keys:
-                if key not in ["title", "link", "snippet", "mime"]:
-                    item.pop(key, None)
-        if results:
-            return results[: int(GPT.SEARCH_NUM_RESULTS)]
-    except Exception as e:
-        logger.error(e)
     return []
src/config.py
@@ -143,9 +143,8 @@ class GPT:
     TEXT_MODEL_NAME = os.getenv("GPT_TEXT_MODEL_NAME", "gpt-4o")  # custom name
     IMAGE_MODEL_NAME = os.getenv("GPT_IMAGE_MODEL_NAME", "gpt-4o")
     VIDEO_MODEL_NAME = os.getenv("GPT_VIDEO_MODEL_NAME", "glm-4v-plus")
-    SEARCH_API_KEY = os.getenv("GPT_SEARCH_API_KEY", "")  # online search (currently, we use GLM)
-    SEARCH_BASE_URL = os.getenv("GPT_SEARCH_BASE_URL", "https://open.bigmodel.cn/api/paas/v4")
-    SEARCH_MODEL = os.getenv("GPT_SEARCH_MODEL", "web-search-pro")
+    GLM_API_KEY = os.getenv("GPT_GLM_API_KEY", "")
+    GLM_BASE_URL = os.getenv("GPT_GLM_BASE_URL", "https://open.bigmodel.cn/api/paas/v4")
     SEARCH_NUM_RESULTS = os.getenv("GPT_SEARCH_NUM_RESULTS", "5")
     PRIMARY_SEARCH_ENGINE = os.getenv("GPT_PRIMARY_SEARCH_ENGINE", "google")  # google or glm
     TEXT_TIMEOUT = os.getenv("GPT_TEXT_TIMEOUT", "120")