flow.py

from pocketflow import Flow
from nodes import (
    ResearcherNode, TodoNode, SearcherNode, 
    EvidenceNode, ReporterNode, GenerateHTMLNode
)
from context import FileSystemContextManager

class OptimizedResearchFlow(Flow):
    def __init__(self):
        super().__init__()
        self.context_manager = FileSystemContextManager()
        self.build_flow()
    
    def build_flow(self):
        """Build the optimized research flow."""
        # Create nodes
        researcher = ResearcherNode()
        todo = TodoNode()
        searcher = SearcherNode()
        evidence = EvidenceNode()
        reporter = ReporterNode()
        generate_html = GenerateHTMLNode()
        
        # Build flow connections
        self.start(researcher)
        researcher.next(todo, "todo")
        
        # Todo node transitions
        todo.next(searcher, "search")
        todo.next(evidence, "evidence")
        todo.next(reporter, "report")
        todo.next(todo, "todo")  # For task progression
        
        # Searcher node transitions
        searcher.next(todo, "todo")
        
        # Evidence node transitions
        evidence.next(reporter, "report")
        
        # Reporter node transitions
        reporter.next(generate_html, "generate_html")
        
        # Generate HTML node transitions
        generate_html.next(None, "done")
    
    def prep(self, shared):
        """Prepare the flow with context manager."""
        # Add context manager to shared
        shared["context_manager"] = self.context_manager
        shared["output_dir"] = self.context_manager.get_session_directory()
        
        # Log start of research
        self.context_manager.save_log(f"Starting research for query: {shared['question']}")
        
        print(f"🚀 Starting optimized research flow...")
        print(f"📁 Session directory: {shared['output_dir']}")
        
        return shared
    
    def post(self, shared, prep_res, exec_res):
        """Finalize the flow and clean up."""
        # Log completion
        self.context_manager.save_log("Research completed successfully")
        
        # Save final state
        final_state = {
            "question": shared.get("question"),
            "report_generated": "report" in shared,
            "session_summary": self.context_manager.get_summary()
        }
        self.context_manager.save_session_state(final_state)
        
        print(f"🎉 Research flow completed!")
        print(f"📊 Session summary: {final_state}")
        
        return exec_res

def create_optimized_flow():
    """Create an instance of the optimized research flow."""
    return OptimizedResearchFlow()