Update article_generator.py

article_generator.py

@@ -47,6 +47,23 @@ class EnhancedTavilySearchTool:
         else:
             raise Exception(f"Failed to fetch data from Tavily API: {response.status_code}, {response.text}")
 
+# ドメインフィルタリング関数の定義
+def filter_responses_by_domain(responses, allowed_domains):
+    filtered_responses = []
+    for response in responses:
+        url = response.get('url')
+        if any(domain in url for domain in allowed_domains):
+            filtered_responses.append(response)
+    return filtered_responses
+
+# 許可されたドメインリスト
+allowed_domains = [
+    '.gov',  # 政府関連のサイト
+    '.edu',  # 教育機関のサイト
+    '.org',  # 非営利組織のサイト
+    '.co.jp'  # 日本の企業サイト
+]
+
 # 重複を排除するヘルパー関数
 def remove_duplicates(text_list):
     seen = set()
@@ -227,10 +244,14 @@ def perform_initial_tavily_search(h2_texts, h3_texts):
     
     print("Performing Tavily search with queries:", queries)
     responses = tavily_search_tool.search(queries)
+
+    # フィルタリングを適用
+    filtered_responses = filter_responses_by_domain(responses, allowed_domains)
+    
     response_dict = {}
     for i, query in enumerate(queries):
-        if i < len(responses):  # 応答リストの範囲内にあることを確認
-            response_dict[query] = responses[i]
+        if i < len(filtered_responses):  # 応答リストの範囲内にあることを確認
+            response_dict[query] = filtered_responses[i]
         else:
             response_dict[query] = "No response received"
     
@@ -440,4 +461,4 @@ def setup_gradio_interface():
                 updated_content = display_content(latest_content.get(), format_choice)
                 return updated_content
 
-            format_selector.change(update_content, inputs=[format_selector], outputs=[content_display])
+            format_selector.change(update_content, inputs=[format_selector], outputs=[content_display])