Update article_generator.py

article_generator.py

@@ -28,14 +28,15 @@ class GoogleSearchTool:
 
 # Tavily APIのカスタムツールを定義
 class EnhancedTavilySearchTool:
-    def search(self, query):
-        if len(query) < 5:
-            query += " details"  # クエリを拡張して必要な文字数にする
+    def search(self, queries):
+        combined_query = " | ".join(queries)  # クエリを結合して一つのリクエストで処理
+        if len(combined_query) < 5:
+            combined_query += " details"
         
         params = {
             'api_key': tavily_api_key,
-            'query': query,
-            'max_results': 5,
+            'query': combined_query,
+            'max_results': 50,  # 必要に応じて結果の数を調整
             'detail_level': 'high',
             'search_depth': 'advanced'
         }
@@ -74,8 +75,8 @@ def clear_state():
     research_results = []
     return "状態がクリアされました"
 
-# 見出しを処理する関数（キャッシュされたレスポンスを使用）
-def process_heading(agent, h2_text, h3_for_this_h2, cached_responses):
+# 見出しを処理する関数
+def process_heading(agent, h2_text, h3_for_this_h2, executed_instructions, research_results, cached_responses):
     query = f"{h2_text} {' '.join(h3_for_this_h2)}"
     if query in cached_responses:
         return (query, cached_responses[query])
@@ -85,26 +86,15 @@ def process_heading(agent, h2_text, h3_for_this_h2, cached_responses):
 # 初期データをTavily検索で収集する関数
 def perform_initial_tavily_search(h2_texts, h3_texts):
     tavily_search_tool = EnhancedTavilySearchTool()
-    cached_responses = {}
     queries = []
-
+    
     for h2_text in h2_texts:
         h3_for_this_h2 = [h3 for h3 in h3_texts if h3.startswith(f"{h2_texts.index(h2_text)+1}-")]
         query = f"{h2_text} {' '.join(h3_for_this_h2)}"
         queries.append(query)
-
-    with ThreadPoolExecutor(max_workers=10) as executor:
-        futures = {executor.submit(tavily_search_tool.search, query): query for query in queries}
-        for future in as_completed(futures):
-            query = futures[future]
-            try:
-                response = future.result()
-                cached_responses[query] = response
-            except Exception as e:
-                print(f"Error occurred during Tavily search for query '{query}': {str(e)}")
-                cached_responses[query] = str(e)
-
-    return cached_responses
+    
+    response = tavily_search_tool.search(queries)  # 一回のリクエストで全てのクエリを処理
+    return {query: response[i] for i, query in enumerate(queries)}  # 結果を適切にマッピング
 
 # キャッシュされたTavilyデータを保存する関数
 def save_preloaded_tavily_data(data):
@@ -159,7 +149,7 @@ def generate_article(editable_output2):
         futures = []
         for h2_text in h2_texts:
             h3_for_this_h2 = [h3 for h3 in h3_texts if h3.startswith(f"{h2_texts.index(h2_text)+1}-")]
-            futures.append(executor.submit(process_heading, agent, h2_text, h3_for_this_h2, cached_responses))
+            futures.append(executor.submit(process_heading, agent, h2_text, h3_for_this_h2, executed_instructions, research_results, cached_responses))
 
         for future in as_completed(futures):
             purpose, response = future.result()
@@ -311,3 +301,4 @@ def continue_generate_article():
         os.remove("state.json")
 
     return final_result
+