Commit d4cfff1
Changed files (2)
src
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")