From fc83d9387e449e8f0509962e16b8ecc815a69807 Mon Sep 17 00:00:00 2001
From: AnotiaWang <anotia0202@gmail.com>
Date: Mon, 17 Feb 2025 16:20:25 +0800
Subject: [PATCH] feat: support setting concurrency limits in web search

---
 README.md                              |  5 +++-
 README_zh.md                           |  2 +-
 components/ConfigManager.vue           | 13 +++++++++
 components/DeepResearch/SearchNode.vue |  1 -
 composables/usePLimit.ts               | 30 ++++++++++++++++++++
 i18n/en.json                           |  4 ++-
 i18n/zh.json                           |  4 ++-
 lib/deep-research.ts                   | 38 +++++++++++++++-----------
 stores/config.ts                       |  5 +++-
 9 files changed, 80 insertions(+), 22 deletions(-)
 create mode 100644 composables/usePLimit.ts

diff --git a/README.md b/README.md
index eb9076a..0845067 100644
--- a/README.md
+++ b/README.md
@@ -24,6 +24,10 @@ Please give a 🌟 Star if you like this project!
 
 ## Recent updates
 
+25/02/17
+
+- Added: set rate limits for web search
+
 25/02/16
 
 - Refactored the search visualization using VueFlow
@@ -48,7 +52,6 @@ Please give a 🌟 Star if you like this project!
 - Added Docker support
 - Fixed "export as PDF" issues
 
-
 ## How to use
 
 Live demo: <a href="https://deep-research.ataw.top" target="_blank">https://deep-research.ataw.top</a>
diff --git a/README_zh.md b/README_zh.md
index bd33b2d..94b9b0f 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -25,7 +25,7 @@
 
 25/02/15
 
-- AI 提供商支持 DeepSeek，OpenRouter 和 Ollama，联网搜素支持 Firecrawl
+- AI 提供商支持 DeepSeek，OpenRouter 和 Ollama，联网搜索支持 Firecrawl
 - 支持检查项目更新
 - 支持重新生成报告
 - 一般性优化和改进
diff --git a/components/ConfigManager.vue b/components/ConfigManager.vue
index 8d821a5..abbdbe0 100644
--- a/components/ConfigManager.vue
+++ b/components/ConfigManager.vue
@@ -239,6 +239,19 @@
               private
             />
           </UFormField>
+          <UFormField :label="$t('settings.webSearch.concurrencyLimit')">
+            <template #help>
+              {{ $t('settings.webSearch.concurrencyLimitHelp') }}
+            </template>
+            <UInput
+              v-model="config.webSearch.concurrencyLimit"
+              class="w-15"
+              type="number"
+              :min="1"
+              :max="5"
+              :step="1"
+            />
+          </UFormField>
         </div>
       </template>
       <template #footer>
diff --git a/components/DeepResearch/SearchNode.vue b/components/DeepResearch/SearchNode.vue
index 7cacda6..1318d37 100644
--- a/components/DeepResearch/SearchNode.vue
+++ b/components/DeepResearch/SearchNode.vue
@@ -1,5 +1,4 @@
 <script setup lang="ts">
-  import { computed, defineProps } from 'vue'
   import { Handle, Position } from '@vue-flow/core'
   import type { ButtonProps } from '@nuxt/ui'
   import type { SearchNodeData } from './SearchFlow.vue'
diff --git a/composables/usePLimit.ts b/composables/usePLimit.ts
new file mode 100644
index 0000000..d425393
--- /dev/null
+++ b/composables/usePLimit.ts
@@ -0,0 +1,30 @@
+import pLimit from 'p-limit'
+
+/**
+ * The concurrency value used by the global limit.
+ * This represents the *actual* limit value.
+ * The value in `globalLimit` should not be used, because `deepResearch` uses recursive calls,
+ * and `globalLimit.concurrency` can be much higher than the actual one.
+ */
+let globalLimitConcurrency = 2
+const globalLimit = pLimit(globalLimitConcurrency)
+
+export function usePLimit() {
+  const { config } = useConfigStore()
+
+  if (
+    config.webSearch.concurrencyLimit &&
+    config.webSearch.concurrencyLimit >= 1 &&
+    globalLimitConcurrency !== config.webSearch.concurrencyLimit
+  ) {
+    console.log(
+      `[usePLimit] Updating concurrency from ${globalLimitConcurrency} to ${config.webSearch.concurrencyLimit}. Current concurrency: ${globalLimit.concurrency}`,
+    )
+    let newLimit = config.webSearch.concurrencyLimit
+    let diff = newLimit - globalLimitConcurrency
+    globalLimitConcurrency = newLimit
+    globalLimit.concurrency += diff
+  }
+
+  return globalLimit
+}
diff --git a/i18n/en.json b/i18n/en.json
index b044a63..322821b 100644
--- a/i18n/en.json
+++ b/i18n/en.json
@@ -38,7 +38,9 @@
         "firecrawl": {
           "help": "Get one API key at {0}."
         }
-      }
+      },
+      "concurrencyLimitHelp": "Limit the concurrent search tasks. This is useful to avoid overloading the search provider and causing requests to fail.",
+      "concurrencyLimit": "Concurrency Limit"
     }
   },
   "researchTopic": {
diff --git a/i18n/zh.json b/i18n/zh.json
index 778226b..4d4d954 100644
--- a/i18n/zh.json
+++ b/i18n/zh.json
@@ -38,7 +38,9 @@
         "tavily": {
           "help": "和 Firecrawl 类似，不过提供了每月 1000 次免费搜索。在 {0} 获取一个 API key。"
         }
-      }
+      },
+      "concurrencyLimit": "并发数",
+      "concurrencyLimitHelp": "限制同时进行的搜索数量。这样可以避免被 API 服务限流，导致请求失败。"
     }
   },
   "researchTopic": {
diff --git a/lib/deep-research.ts b/lib/deep-research.ts
index e56e007..7c6dbda 100644
--- a/lib/deep-research.ts
+++ b/lib/deep-research.ts
@@ -61,9 +61,6 @@ export type ResearchStep =
   | { type: 'error'; message: string; nodeId: string }
   | { type: 'complete'; learnings: string[]; visitedUrls: string[] }
 
-// increase this if you have higher API rate limits
-const ConcurrencyLimit = 2
-
 /**
  * Schema for {@link generateSearchQueries} without dynamic descriptions
  */
@@ -242,6 +239,7 @@ export async function deepResearch({
 }): Promise<ResearchResult> {
   const { t } = useNuxtApp().$i18n
   const language = t('language', {}, { locale: languageCode })
+  const globalLimit = usePLimit()
 
   onProgress({
     type: 'generating_query',
@@ -257,7 +255,6 @@ export async function deepResearch({
       language,
       searchLanguage,
     })
-    const limit = pLimit(ConcurrencyLimit)
 
     let searchQueries: PartialSearchQuery[] = []
 
@@ -322,7 +319,7 @@ export async function deepResearch({
     // Run in parallel and limit the concurrency
     const results = await Promise.all(
       searchQueries.map((searchQuery, i) =>
-        limit(async () => {
+        globalLimit(async () => {
           if (!searchQuery?.query) {
             return {
               learnings: [],
@@ -433,17 +430,26 @@ export async function deepResearch({
                 .join('')}
             `.trim()
 
-              return deepResearch({
-                query: nextQuery,
-                breadth: nextBreadth,
-                maxDepth,
-                learnings: allLearnings,
-                visitedUrls: allUrls,
-                onProgress,
-                currentDepth: nextDepth,
-                nodeId: childNodeId(nodeId, i),
-                languageCode,
-              })
+              // Add concurrency by 1, and do next recursive search
+              globalLimit.concurrency++
+              try {
+                const r = await deepResearch({
+                  query: nextQuery,
+                  breadth: nextBreadth,
+                  maxDepth,
+                  learnings: allLearnings,
+                  visitedUrls: allUrls,
+                  onProgress,
+                  currentDepth: nextDepth,
+                  nodeId: childNodeId(nodeId, i),
+                  languageCode,
+                })
+                return r
+              } catch (error) {
+                throw error
+              } finally {
+                globalLimit.concurrency--
+              }
             } else {
               return {
                 learnings: allLearnings,
diff --git a/stores/config.ts b/stores/config.ts
index 9fa70bc..1192029 100644
--- a/stores/config.ts
+++ b/stores/config.ts
@@ -21,6 +21,8 @@ export interface ConfigWebSearch {
   apiKey?: string
   /** Force the LLM to generate serp queries in a certain language */
   searchLanguage?: Locale
+  /** Limit the number of concurrent tasks globally */
+  concurrencyLimit?: number
 }
 
 export interface Config {
@@ -37,8 +39,9 @@ export const useConfigStore = defineStore('config', () => {
     },
     webSearch: {
       provider: 'tavily',
+      concurrencyLimit: 2,
     },
-  })
+  } satisfies Config)
   // The version user dismissed the update notification
   const dismissUpdateVersion = useLocalStorage<string>(
     'dismiss-update-version',