Commit 041eccf

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-02-10 15:07:57
feat(gpt): customize the number of search results
1 parent 71deba1
Changed files (2)
src/llm/tool_call.py
@@ -34,7 +34,7 @@ async def get_online_search_result(query: str) -> list[dict]:
             x.pop("index", None)
             x.pop("refer", None)
         if results:
-            return results
+            return results[: int(GPT.SEARCH_NUM_RESULTS)]
     except Exception as e:
         logger.error(e)
     return await google_search_result(query)
@@ -51,7 +51,7 @@ async def google_search_result(query: str) -> list[dict]:
                 if key not in ["title", "link", "snippet", "mime"]:
                     item.pop(key, None)
         if results:
-            return results[:5]
+            return results[: int(GPT.SEARCH_NUM_RESULTS)]
     except Exception as e:
         logger.error(e)
     return []
src/config.py
@@ -146,6 +146,7 @@ class GPT:
     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")
+    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")
     IMAGE_TIMEOUT = os.getenv("GPT_IMAGE_TIMEOUT", "120")