Update research_agent/agent.py

research_agent/agent.py

@@ -80,15 +80,23 @@ def write_report_stream(config: AgentConfig, writer_model, tavily_client, embedd
     detailed_topic = plan["detailed_topic"]
     sections = plan["sections"]
     
-    yield f"### Starting Report Generation for: {detailed_topic}\n\n"
+    yield f"### Starting Report Generation for: {detailed_topic}\n\n---\n"
     
-    report_state = {"full_report_text": f"# Deep Research Report: {detailed_topic}\n\n", "all_source_urls": set()}
+    report_state = {
+        "full_report_text": f"# Deep Research Report: {detailed_topic}\n\n", 
+        "all_source_urls": set(),
+        "completed_section_titles": []
+    }
     rag_pipeline = RAGPipeline(embedding_model, reranker)
     
     for i, section in enumerate(sections):
-        yield f"--- \n### Processing Section {i+1}/{len(sections)}: {section.title}...\n"
+        yield f"### Processing Section {i+1}/{len(sections)}: {section.title}...\n\n"
         
-        previous_sections_context = report_state["full_report_text"]
+        previous_sections_context = "\n".join(f"- {title}" for title in report_state["completed_section_titles"])
+        if previous_sections_context:
+            previous_sections_context = "The following sections have already been written:\n" + previous_sections_context
+        else:
+            previous_sections_context = "This is the first section."
         
         section_queries = [f"{detailed_topic} - {section.title}"] + section.description.split('\n')[-3:]
         section_queries = [q.strip() for q in section_queries if q.strip()]
@@ -121,6 +129,7 @@ def write_report_stream(config: AgentConfig, writer_model, tavily_client, embedd
         bibliography = "\n".join(f"[{num}] {url}" for url, num in cited_sources_for_section.items())
 
         yield f"-> Synthesizing and writing section content...\n"
+        yield f"--> Calling LLM for section '{section.title}'. This may take a moment...\n"
         writer_prompt = prompts.writer_prompt_template.format(
             writer_system_instruction=prompts.writer_system_instruction,
             previous_sections_context=previous_sections_context,
@@ -130,14 +139,14 @@ def write_report_stream(config: AgentConfig, writer_model, tavily_client, embedd
         
         draft_content = run_gemini_text_completion(writer_model, writer_prompt, config.WRITER_TEMPERATURE)
         
-        yield "-> Fact-checking and verifying section...\n"
+        yield "-> Fact-checking and verifying section...\n\n---\n"
         final_content = run_verification_step(writer_model, draft_content, context_for_llm)
         final_content_with_sources = f"## {section.title}\n\n{final_content}\n\n**Sources Used in this Section**\n{bibliography}\n\n"
         
         report_state["full_report_text"] += final_content_with_sources
+        report_state["completed_section_titles"].append(section.title)
 
     final_bibliography = "\n".join(f"- {url}" for url in sorted(list(report_state["all_source_urls"])))
     report_state["full_report_text"] += f"## Master Bibliography\n\n{final_bibliography}"
 
-    yield "\n--- Report Generation Complete ---\n"
     return report_state["full_report_text"]