Commit 5f19ecd

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-02-10 13:25:39
feat(gpt): add google search
1 parent 3d3fdb8
Changed files (3)
src/llm/response.py
@@ -67,9 +67,10 @@ async def parse_tool_call(config: dict, response: dict, retry: int = 0, **kwargs
     Returns:
         {"content": str, "reasoning": str}
     """
-    choice = glom(response, "choices.0", default=[])
-    if not choice:
-        return {"content": "", "reasoning": ""}
+    choice = glom(response, "choices.0", default={})
+    if choice.get("finish_reason", "") not in ["stop", "tool_calls", "length"]:
+        logger.warning(response)
+        raise  # noqa: PLE0704
     logger.debug(response)
     try:
         if tool_call := glom(choice, "message.tool_calls.0", default={}):
src/llm/tool_call.py
@@ -1,10 +1,13 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+import copy
+
 from glom import glom
 from loguru import logger
 from openai import AsyncOpenAI, DefaultAsyncHttpxClient
 
-from config import GPT, PROXY
+from config import GPT, PROXY, TOKEN
+from networking import hx_req
 
 
 async def get_online_search_result(query: str) -> list[dict]:
@@ -28,7 +31,25 @@ async def get_online_search_result(query: str) -> list[dict]:
             x.pop("icon", None)
             x.pop("index", None)
             x.pop("refer", None)
-        return results
+        if results:
+            return 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[:5]
     except Exception as e:
         logger.error(e)
-        return []
+    return []
src/config.py
@@ -109,6 +109,8 @@ class TOKEN:
     TENCENT_ASR_SECRET_KEY = os.getenv("TENCENT_ASR_SECRET_KEY", "")
     YOUTUBE_API_KEY = os.getenv("YOUTUBE_API_KEY", "")
     CMC_API_KEY = os.getenv("CMC_API_KEY", "")
+    GOOGLE_SEARCH_API_KEY = os.getenv("GOOGLE_SEARCH_API_KEY", "")
+    GOOGLE_SEARCH_CX = os.getenv("GOOGLE_SEARCH_CX", "")
 
 
 class PROXY:  # format: socks5://127.0.0.1:7890
@@ -119,6 +121,7 @@ class PROXY:  # format: socks5://127.0.0.1:7890
     GPT = os.getenv("GPT_PROXY", None)
     SUBTITLE = os.getenv("SUBTITLE_PROXY", None)
     CRYPTO = os.getenv("CRYPTO_PROXY", None)
+    GOOGLE_SEARCH = os.getenv("GOOGLE_SEARCH_PROXY", None)
     DOWNLOAD = os.getenv("DOWNLOAD_PROXY", None)
     WEIBO_COOKIE = os.getenv("WEIBO_COOKIE_PROXY", None)  # Weibo visitor cookie
     YTDLP = os.getenv("YTDLP_PROXY", None)  # general proxy for ytdlp