Update article_generator.py

article_generator.py

@@ -63,7 +63,7 @@ def save_state(state):
 # 状態をロードする関数
 def load_state():
     if os.path.exists(state_file):
-        with open(state_file, "r", encoding="utf-8") as f:
+        with open(state_file, "r", encoding="utf-8") as f):
             return json.load(f)
     return None
 
@@ -125,6 +125,53 @@ def setup_plan_and_execute_agent():
     agent = PlanAndExecute(planner=planner, executor=executor, verbose=True)
     return agent
 
+# 記事のセクションをGPT-4で拡張する関数
+def expand_section_with_gpt4(h2_text, h3_texts, preloaded_data):
+    prompts = []
+    for h3_text in h3_texts:
+        key = f"{h2_text} {h3_text}"
+        if key in preloaded_data:
+            context = preloaded_data[key]
+            prompt = f"このテーマ「{h3_text}」について、さらに詳細な情報を加えてください。以下は初期の情報です：\n{context}"
+            prompts.append(prompt)
+        else:
+            prompts.append(f"このテーマ「{h3_text}」に関する詳細情報を加えてください。")
+    
+    expanded_texts = []
+    for prompt in prompts:
+        try:
+            response = openai.Completion.create(
+                model="gpt-4-turbo",
+                prompt=prompt,
+                max_tokens=2000,
+                temperature=0.7
+            )
+            expanded_text = response.choices[0].text.strip()
+            expanded_texts.append(expanded_text)
+        except Exception as e:
+            print(f"Error in generating text for {prompt}: {str(e)}")
+            expanded_texts.append("Error in text generation.")
+    
+    return expanded_texts
+
+# 記事を拡張する関数
+def generate_expanded_article(article_html):
+    soup = BeautifulSoup(article_html, 'html.parser')
+    h2_elements = soup.find_all('h2')
+    preloaded_data = load_preloaded_tavily_data()
+    
+    for h2 in h2_elements:
+        h3_texts = [h3.get_text() for h3 in h2.find_next_siblings('h3', limit=2)]
+        expanded_texts = expand_section_with_gpt4(h2.get_text(), h3_texts, preloaded_data)
+        
+        for h3, expanded_text in zip(h3_texts, expanded_texts):
+            h3_tag = soup.find('h3', text=h3)
+            new_paragraph = soup.new_tag('p')
+            new_paragraph.string = expanded_text
+            h3_tag.insert_after(new_paragraph)
+    
+    return str(soup)
+
 # 記事を生成する関数
 def generate_article(editable_output2):
     # 途中から再開する場合のために状態を読み込み
@@ -237,16 +284,19 @@ def generate_article(editable_output2):
 
     final_result = "\n".join(results)
 
+    # 生成された初期記事を拡張
+    expanded_article = generate_expanded_article(final_result)
+
     with open("output3.txt", "w", encoding="utf-8") as f:
-        f.write(final_result)
+        f.write(expanded_article)
 
-    print(final_result)  # ログに最終結果を出力
+    print(expanded_article)  # ログに最終結果を出力
 
     # 生成が完了したら状態ファイルを削除
     if os.path.exists("state.json"):
         os.remove("state.json")
 
-    return final_result
+    return expanded_article
 
 def continue_generate_article():
     state = load_state()
@@ -301,4 +351,4 @@ def continue_generate_article():
     if os.path.exists("state.json"):
         os.remove("state.json")
 
-    return final_result
+    return final_result