Update article_generator.py

article_generator.py

@@ -211,30 +211,47 @@ def generate_text_with_gpt4(prompt):
     )
     return response.choices[0]["message"]["content"].strip()
 
-# 初期データをTavily検索で収集する関数
+# Tavily APIからデータを取得し、結果とURLを保存する関数
 def perform_initial_tavily_search(h2_texts, h3_texts):
     tavily_search_tool = EnhancedTavilySearchTool()
     queries = []
-    
+    references = {}  # URL参照を保存するための辞書
+
     for idx, h2_text in enumerate(h2_texts):
         h3_for_this_h2 = [h3 for h3 in h3_texts if h3.startswith(f"{idx+1}-")]
         if not h3_for_this_h2 and h2_text.strip() != "まとめ":  # "まとめ" セクションを除外
-            print(f"No matching h3 elements found for h2: {h2_text} at index {idx+1}")
             continue
-        
+
         query = f"{h2_text} {' '.join(h3_for_this_h2)}"
         queries.append(query)
     
-    print("Performing Tavily search with queries:", queries)
     responses = tavily_search_tool.search(queries)
     response_dict = {}
     for i, query in enumerate(queries):
-        if i < len(responses):  # 応答リストの範囲内にあることを確認
+        if i < len(responses):
             response_dict[query] = responses[i]
+            references[query] = responses[i].get('url', 'No URL available')  # URLを保存
         else:
             response_dict[query] = "No response received"
-    
-    return response_dict
+
+    return response_dict, references
+
+# 記事の最後に参照URLを追加する関数
+def add_references_section(soup, references):
+    ref_section = soup.new_tag('div', **{'class': 'references'})
+    ref_title = soup.new_tag('h2')
+    ref_title.string = "参照"
+    ref_section.append(ref_title)
+
+    for ref in references.values():
+        ref_tag = soup.new_tag('p')
+        ref_link = soup.new_tag('a', href=ref)
+        ref_link.string = ref
+        ref_tag.append(ref_link)
+        ref_section.append(ref_tag)
+
+    soup.append(ref_section)
+    return soup
 
 def save_preloaded_tavily_data(data):
     with open("preloaded_tavily_data.json", "w", encoding="utf-8") as f:
@@ -268,7 +285,7 @@ def generate_article(editable_output2):
 
     # 初期のTavily検索
     print("Performing initial Tavily search...")
-    cached_responses = perform_initial_tavily_search(h2_texts, h3_texts)
+    cached_responses, references = perform_initial_tavily_search(h2_texts, h3_texts)
     save_preloaded_tavily_data(cached_responses)
 
     executed_instructions = []
@@ -369,21 +386,24 @@ def generate_article(editable_output2):
 
     # 初期データをTavily検索で収集する関数
     h3_texts = [h3.get_text(strip=True) for h3 in updated_soup.find_all('h3')]
-    cached_responses = perform_initial_tavily_search([], h3_texts)
+    cached_responses, references = perform_initial_tavily_search([], h3_texts)
     save_preloaded_tavily_data(cached_responses)
 
     # h3タグの拡張を行う
     expanded_soup = expand_h3_sections(updated_soup, cached_responses)
 
     final_html = str(expanded_soup)
-    final_markdown = custom_html_to_markdown(final_html)
+    final_soup_with_refs = add_references_section(BeautifulSoup(final_html, 'html.parser'), references)  # 参照セクションの追加
+    final_html_with_refs = str(final_soup_with_refs)
 
     with open("output3.txt", "w", encoding="utf-8") as f:
-        f.write(final_html)
+        f.write(final_html_with_refs)
 
-    print("Article generation complete. Output saved to output3.txt.")
-    return final_markdown, final_html
+    print("Final generated article content with references:")
+    print(final_html_with_refs)
 
+    print("Article generation complete. Output saved to output3.txt.")
+    return final_html_with_refs
 
 # HTMLをMarkdownに変換する関数
 def custom_html_to_markdown(html):