Update app.py

app.py

@@ -105,7 +105,7 @@ class Task:
         created_at (datetime): The timestamp when the task was created.
         completed_at (Optional[datetime]): The timestamp when the task was completed or failed.
         performance_metrics (Dict[str, float]): Metrics related to the task's performance, like execution time.
-        metadata (Dict[str, Any]): A dictionary for storing arbitrary data related to the task, such as the suggested agent role.
+        context (str): A string containing the results of previous tasks to provide context for the current one.
     """
     id: str
     description: str
@@ -117,7 +117,7 @@ class Task:
     created_at: datetime = field(default_factory=datetime.now)
     completed_at: Optional[datetime] = None
     performance_metrics: Dict[str, float] = field(default_factory=dict)
-    metadata: Dict[str, Any] = field(default_factory=dict)
+    context: str = ""
 
 
 # ==============================================================================
@@ -129,8 +129,8 @@ class BaseAgent:
     An abstract base class for all AI agents in the system.
 
     This class provides the fundamental structure for agents, including task
-    processing logic, memory management, and collaboration protocols. Each
-    specialized agent extends this class to implement role-specific behaviors.
+    processing logic. Each specialized agent extends this class to implement
+    role-specific behaviors.
     """
     def __init__(self, name: str, role: AgentRole, llm: Optional[ChatOpenAI] = None):
         """
@@ -146,7 +146,6 @@ class BaseAgent:
         self.role = role
         self.llm = llm
         self.current_task: Optional[Task] = None
-        self.completed_tasks: List[Task] = []
 
     async def process_task(self, task: Task) -> Task:
         """
@@ -175,7 +174,6 @@ class BaseAgent:
 
             task.result = result
             task.status = TaskStatus.COMPLETED
-            task.completed_at = datetime.now()
             task.confidence = self._calculate_confidence(result)
 
         except Exception as e:
@@ -186,11 +184,9 @@ class BaseAgent:
         
         finally:
             # Record performance metrics regardless of success or failure
-            if task.status != TaskStatus.IN_PROGRESS:
-                task.completed_at = task.completed_at or datetime.now()
-                execution_time = (task.completed_at - start_time).total_seconds()
-                task.performance_metrics['execution_time'] = execution_time
-                self.completed_tasks.append(task)
+            task.completed_at = datetime.now()
+            execution_time = (task.completed_at - start_time).total_seconds()
+            task.performance_metrics['execution_time'] = execution_time
             self.current_task = None
 
         return task
@@ -250,8 +246,8 @@ class AnalystAgent(BaseAgent):
 
     async def _execute_task(self, task: Task) -> Any:
         prompt = ChatPromptTemplate.from_messages([
-            SystemMessage(content="You are an expert analyst. Your role is to dissect information, identify underlying patterns, trends, and correlations, and present actionable insights. Your analysis must be logical and evidence-based."),
-            HumanMessage(content=f"Please analyze the following information and provide a detailed breakdown: {task.description}")
+            SystemMessage(content="You are an expert analyst. Your role is to dissect information, identify underlying patterns, trends, and correlations, and present actionable insights. Use the provided context as the basis for your analysis."),
+            HumanMessage(content=f"Based on the following context, please perform your analysis.\n\nCONTEXT:\n{task.context}\n\nTASK:\n{task.description}")
         ])
         response = await self.llm.ainvoke(prompt.format_messages())
         return {"analysis": response.content, "patterns": "Identified key performance indicators and trends."}
@@ -263,8 +259,8 @@ class CriticAgent(BaseAgent):
 
     async def _execute_task(self, task: Task) -> Any:
         prompt = ChatPromptTemplate.from_messages([
-            SystemMessage(content="You are a meticulous critic. Your function is to rigorously evaluate the provided information, identifying logical fallacies, biases, assumptions, and gaps. Provide constructive feedback for improvement."),
-            HumanMessage(content=f"Please provide a critical evaluation of the following: {task.description}")
+            SystemMessage(content="You are a meticulous critic. Your function is to rigorously evaluate the provided information, identifying logical fallacies, biases, and gaps. Provide constructive feedback for improvement. Base your critique on the context provided."),
+            HumanMessage(content=f"Based on the following context, please perform your critique.\n\nCONTEXT:\n{task.context}\n\nTASK:\n{task.description}")
         ])
         response = await self.llm.ainvoke(prompt.format_messages())
         return {"evaluation": response.content, "strengths": "Identified robust arguments.", "weaknesses": "Flagged potential biases."}
@@ -276,8 +272,8 @@ class SynthesizerAgent(BaseAgent):
 
     async def _execute_task(self, task: Task) -> Any:
         prompt = ChatPromptTemplate.from_messages([
-            SystemMessage(content="You are an expert synthesizer. Your task is to integrate disparate pieces of information, resolve contradictions, and formulate a single, coherent, and comprehensive narrative or strategic plan."),
-            HumanMessage(content=f"Please synthesize the following inputs into a unified conclusion: {task.description}")
+            SystemMessage(content="You are an expert synthesizer. Your task is to integrate all the provided information from the context (research, analysis, critique) into a single, coherent, and comprehensive final report with actionable recommendations."),
+            HumanMessage(content=f"Please synthesize the following inputs from the context into a unified conclusion.\n\nCONTEXT:\n{task.context}\n\nTASK:\n{task.description}")
         ])
         response = await self.llm.ainvoke(prompt.format_messages())
         return {"synthesis": response.content, "recommendations": "Formulated final strategic recommendations."}
@@ -292,8 +288,8 @@ class CoordinatorAgent(BaseAgent):
     The central agent responsible for managing the entire workflow.
 
     The Coordinator decomposes the main problem, creates and assigns tasks,
-    builds a dependency graph, and orchestrates the execution of the workflow
-    by the specialized agents.
+    builds a dependency graph, and orchestrates the execution of the workflow,
+    ensuring information flows correctly between agents.
     """
     def __init__(self, name: str, llm: Optional[ChatOpenAI] = None):
         super().__init__(name, AgentRole.COORDINATOR, llm)
@@ -317,12 +313,11 @@ class CoordinatorAgent(BaseAgent):
         Returns:
             List[Task]: A list of Task objects ready for execution.
         """
-        # A standard workflow template ensures consistency.
         tasks = [
-            Task(id="task_1", description=f"Conduct foundational research on: {problem}", metadata={"suggested_role": "Researcher"}),
-            Task(id="task_2", description=f"Analyze the research findings for patterns and insights related to: {problem}", metadata={"suggested_role": "Analyst"}),
-            Task(id="task_3", description="Critically evaluate the research and analysis for quality, bias, and completeness.", metadata={"suggested_role": "Critic"}),
-            Task(id="task_4", description="Synthesize all findings into a final report with actionable recommendations.", metadata={"suggested_role": "Synthesizer"})
+            Task(id="task_1", description=f"Conduct foundational research on: {problem}"),
+            Task(id="task_2", description=f"Analyze the research findings for patterns and insights related to: {problem}"),
+            Task(id="task_3", description="Critically evaluate the research and analysis for quality, bias, and completeness."),
+            Task(id="task_4", description="Synthesize all findings into a final report with actionable recommendations.")
         ]
         self._build_dependency_graph(tasks)
         return tasks
@@ -334,10 +329,8 @@ class CoordinatorAgent(BaseAgent):
 
     async def execute_workflow(self, tasks: List[Task]) -> Dict[str, Any]:
         """
-        Executes a list of tasks according to their dependencies.
-
-        This method uses a thread pool to execute tasks in parallel where possible,
-        respecting the predefined dependency graph.
+        Executes a list of tasks according to their dependencies, managing the
+        flow of information between them.
 
         Args:
             tasks (List[Task]): The list of tasks to execute.
@@ -349,6 +342,7 @@ class CoordinatorAgent(BaseAgent):
         self._update_workflow_graph_with_tasks(tasks)
 
         completed_task_ids = set()
+        task_results = {}
         task_dict = {t.id: t for t in tasks}
 
         while len(completed_task_ids) < len(tasks):
@@ -356,45 +350,48 @@ class CoordinatorAgent(BaseAgent):
                 t for t in tasks if t.status == TaskStatus.PENDING and all(dep in completed_task_ids for dep in t.dependencies)
             ]
             if not ready_tasks:
-                # Break if no tasks are ready to run to prevent infinite loops
-                failed_tasks = [t for t in tasks if t.status == TaskStatus.FAILED]
-                if len(completed_task_ids) + len(failed_tasks) == len(tasks):
+                if any(t.status == TaskStatus.FAILED for t in tasks):
+                    logger.warning("Workflow halted due to a failed task.")
                     break
-                # If there are still pending tasks but none are ready, it indicates a dependency issue.
-                logger.error("Workflow stalled: circular dependency or unresolved failed dependency.")
+                logger.error("Workflow stalled: No ready tasks found. Check for circular dependencies.")
                 break
 
+            # Prepare context for the ready tasks
+            for task in ready_tasks:
+                context_parts = []
+                for dep_id in task.dependencies:
+                    if dep_id in task_results:
+                        context_parts.append(f"--- Result of {dep_id} ---\n{json.dumps(task_results[dep_id], indent=2)}")
+                task.context = "\n\n".join(context_parts)
+
             with ThreadPoolExecutor(max_workers=len(self.agents)) as executor:
-                future_to_task = {}
-                for task in ready_tasks:
-                    agent = self._select_agent_for_task(task)
-                    if agent:
-                        task.status = TaskStatus.IN_PROGRESS
-                        future = executor.submit(asyncio.run, agent.process_task(task))
-                        future_to_task[future] = task.id
+                future_to_task = {executor.submit(asyncio.run, self._select_agent_for_task(t).process_task(t)): t for t in ready_tasks}
 
                 for future in as_completed(future_to_task):
-                    task_id = future_to_task[future]
-                    try:
-                        completed_task = future.result()
-                        task_dict[task_id] = completed_task
-                        if completed_task.status == TaskStatus.COMPLETED:
-                            completed_task_ids.add(task_id)
-                        # Update the graph with the final status
-                        self.workflow_graph.nodes[task_id]['status'] = completed_task.status.value
-                    except Exception as exc:
-                        logger.error(f"Task {task_id} generated an exception: {exc}")
-                        task_dict[task_id].status = TaskStatus.FAILED
-                        self.workflow_graph.nodes[task_id]['status'] = TaskStatus.FAILED.value
+                    completed_task = future.result()
+                    task_id = completed_task.id
+                    task_dict[task_id] = completed_task
+                    
+                    if completed_task.status == TaskStatus.COMPLETED:
+                        completed_task_ids.add(task_id)
+                        task_results[task_id] = completed_task.result
+                    
+                    self.workflow_graph.nodes[task_id]['status'] = completed_task.status.value
         
         final_tasks = list(task_dict.values())
         return self._compile_workflow_results(final_tasks, start_time)
 
     def _select_agent_for_task(self, task: Task) -> Optional[BaseAgent]:
         """Selects an available agent best suited for a given task."""
-        suggested_role = task.metadata.get("suggested_role")
+        role_map = {
+            "task_1": AgentRole.RESEARCHER,
+            "task_2": AgentRole.ANALYST,
+            "task_3": AgentRole.CRITIC,
+            "task_4": AgentRole.SYNTHESIZER,
+        }
+        target_role = role_map.get(task.id)
         for agent in self.agents.values():
-            if agent.role.value == suggested_role:
+            if agent.role == target_role:
                 return agent
         return None
 
@@ -404,7 +401,6 @@ class CoordinatorAgent(BaseAgent):
             self.workflow_graph.add_node(task.id, task_description=task.description, status=task.status.value)
             for dep_id in task.dependencies:
                 self.workflow_graph.add_edge(dep_id, task.id)
-            # Link agent to the task it will perform
             agent = self._select_agent_for_task(task)
             if agent:
                 self.workflow_graph.add_edge(agent.name, task.id)
@@ -445,7 +441,6 @@ class WorkflowVisualizer:
         if not G.nodes():
             return self._create_empty_figure("No workflow data available.")
 
-        # Use a hierarchical layout
         pos = nx.spring_layout(G, k=0.9, iterations=50, seed=42)
 
         edge_x, edge_y = [], []
@@ -469,7 +464,7 @@ class WorkflowVisualizer:
                 node_text.append(f"<b>{node}</b><br>{role}")
                 node_colors.append(self.color_map.get(role, '#ccc'))
                 node_sizes.append(35)
-            else: # is task
+            else:
                 status = G.nodes[node].get('status', 'Pending')
                 node_text.append(f"<b>{node}</b><br>Status: {status}")
                 node_colors.append('#6c757d' if status == 'Pending' else '#28a745' if status == 'Completed' else '#dc3545')
@@ -486,13 +481,13 @@ class WorkflowVisualizer:
 
     def create_task_timeline(self, tasks: List[Task]) -> go.Figure:
         """Creates a Plotly timeline (Gantt chart) of task execution."""
-        if not tasks or not any(t.created_at and t.completed_at for t in tasks):
+        if not any(t.created_at and t.completed_at for t in tasks):
             return self._create_empty_figure("No task execution data to display.")
             
-        df_data = []
-        for task in tasks:
-            if task.created_at and task.completed_at:
-                 df_data.append(dict(Task=task.id, Start=task.created_at, Finish=task.completed_at, Agent=task.assigned_to or "Unassigned"))
+        df_data = [
+            dict(Task=task.id, Start=task.created_at, Finish=task.completed_at, Agent=task.assigned_to or "Unassigned")
+            for task in tasks if task.created_at and task.completed_at
+        ]
         
         if not df_data:
             return self._create_empty_figure("No completed tasks with timing data.")
@@ -552,7 +547,7 @@ class ReportGenerator:
             str: A formatted HTML string representing the report.
         """
         tasks = workflow_result.get('tasks', [])
-        synthesis_task = next((t for t in reversed(tasks) if t.status == TaskStatus.COMPLETED and t.metadata.get("suggested_role") == "Synthesizer"), None)
+        synthesis_task = next((t for t in reversed(tasks) if t.status == TaskStatus.COMPLETED and self._get_task_role(t) == AgentRole.SYNTHESIZER), None)
         
         report = f"""
         <div style="font-family: sans-serif; color: #333;">
@@ -568,6 +563,11 @@ class ReportGenerator:
         """
         return report
 
+    def _get_task_role(self, task: Task) -> Optional[AgentRole]:
+        """Helper to determine the role associated with a task."""
+        role_map = { "task_1": AgentRole.RESEARCHER, "task_2": AgentRole.ANALYST, "task_3": AgentRole.CRITIC, "task_4": AgentRole.SYNTHESIZER }
+        return role_map.get(task.id)
+
     def _generate_summary(self, result: Dict[str, Any]) -> str:
         """Generates the executive summary section of the report."""
         return f"""
@@ -579,10 +579,10 @@ class ReportGenerator:
 
     def _generate_recommendations(self, synth_task: Optional[Task]) -> str:
         """Generates the key recommendations section from the Synthesizer's output."""
-        content = "<p>No synthesized recommendations were produced.</p>"
+        content = "<p>No synthesized recommendations were produced due to workflow incompletion or error.</p>"
         if synth_task and isinstance(synth_task.result, dict):
             recommendations = synth_task.result.get('synthesis', 'No specific recommendations provided in the synthesis.')
-            content = f"<p>{recommendations}</p>"
+            content = f"<div>{recommendations}</div>"
 
         return f"""
         <div style="margin-bottom: 20px;">
@@ -635,12 +635,7 @@ def create_gradio_interface():
     visualizer = WorkflowVisualizer()
     report_generator = ReportGenerator()
 
-    # Use a dictionary for shared state to avoid global variables
-    state = {
-        "coordinator": None,
-        "current_workflow": None,
-        "current_problem": ""
-    }
+    state = { "coordinator": None }
 
     def initialize_system(api_key: str, model: str, demo_mode: bool) -> str:
         """Initializes the coordinator and specialized agents."""
@@ -665,19 +660,18 @@ def create_gradio_interface():
 
     async def run_analysis(problem: str) -> Tuple[str, go.Figure, go.Figure, go.Figure, str]:
         """Runs the full analysis workflow for a given problem."""
-        if not state["coordinator"]:
-            return "Error: System not initialized. Please initialize first.", None, None, None, ""
+        coordinator = state.get("coordinator")
+        if not coordinator:
+            empty_fig = visualizer._create_empty_figure("System not initialized.")
+            return "Error: System not initialized.", empty_fig, empty_fig, empty_fig, ""
         if not problem:
-            return "Error: Problem statement cannot be empty.", None, None, None, ""
-        
-        state["current_problem"] = problem
+            empty_fig = visualizer._create_empty_figure("No problem statement provided.")
+            return "Error: Problem statement cannot be empty.", empty_fig, empty_fig, empty_fig, ""
         
         try:
-            tasks = state["coordinator"].decompose_problem(problem)
-            workflow = await state["coordinator"].execute_workflow(tasks)
-            state["current_workflow"] = workflow
+            tasks = coordinator.decompose_problem(problem)
+            workflow = await coordinator.execute_workflow(tasks)
 
-            # Generate outputs
             status_text = f"Analysis complete. Success Rate: {workflow['success_rate']:.0%}. Total Time: {workflow['execution_time']:.1f}s."
             graph_fig = visualizer.create_workflow_graph(workflow['workflow_graph'])
             timeline_fig = visualizer.create_task_timeline(workflow['tasks'])
@@ -686,10 +680,10 @@ def create_gradio_interface():
 
             return status_text, graph_fig, timeline_fig, perf_fig, report_html
         except Exception as e:
-            logger.error(f"An error occurred during analysis: {e}")
-            return f"An unexpected error occurred: {e}", None, None, None, ""
+            logger.error(f"An error occurred during analysis: {e}", exc_info=True)
+            empty_fig = visualizer._create_empty_figure(f"An error occurred: {e}")
+            return f"An unexpected error occurred: {e}", empty_fig, empty_fig, empty_fig, ""
 
-    # Define CSS for a professional look and feel
     custom_css = """
     .gradio-container { max-width: 1400px !important; margin: auto !important; }
     h1 { color: #003366; font-family: sans-serif; text-align: center; }
@@ -702,7 +696,6 @@ def create_gradio_interface():
         
         with gr.Row():
             with gr.Column(scale=1):
-                # Configuration Panel
                 with gr.Accordion("System Configuration", open=True):
                     api_key_input = gr.Textbox(label="OpenAI API Key", type="password", info="Required for live mode.")
                     model_select = gr.Dropdown(choices=["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"], value=Config.DEFAULT_MODEL, label="Language Model")
@@ -711,7 +704,6 @@ def create_gradio_interface():
                     init_status = gr.Textbox(label="System Status", interactive=False)
             
             with gr.Column(scale=3):
-                # Main Analysis Panel
                 with gr.Group():
                     problem_input = gr.Textbox(label="Problem Statement", placeholder="Enter a complex problem for the multi-agent system to analyze...", lines=3)
                     analyze_button = gr.Button("Run Analysis", variant="primary")
@@ -728,17 +720,8 @@ def create_gradio_interface():
             with gr.TabItem("Generated Report"):
                 report_output = gr.HTML()
 
-        # Event Handlers
-        init_button.click(
-            fn=initialize_system,
-            inputs=[api_key_input, model_select, demo_mode_checkbox],
-            outputs=init_status
-        )
-        analyze_button.click(
-            fn=lambda p: asyncio.run(run_analysis(p)),
-            inputs=[problem_input],
-            outputs=[analysis_status, workflow_graph, timeline_chart, performance_chart, report_output]
-        )
+        init_button.click(fn=initialize_system, inputs=[api_key_input, model_select, demo_mode_checkbox], outputs=init_status)
+        analyze_button.click(fn=lambda p: asyncio.run(run_analysis(p)), inputs=[problem_input], outputs=[analysis_status, workflow_graph, timeline_chart, performance_chart, report_output])
 
     return interface