Merge pull request #518 from MervinPraison/claude/issue-484-20250527_043953

feat: Add context management for workflow loops

Author: MervinPraison

src/praisonai-agents/praisonaiagents/process/process.py

@@ -30,6 +30,44 @@ def __init__(self, tasks: Dict[str, Task], agents: List[Agent], manager_llm: Opt
         self.task_retry_counter: Dict[str, int] = {} # Initialize retry counter
         self.workflow_finished = False # ADDED: Workflow finished flag
 
+    def _build_task_context(self, current_task: Task) -> str:
+        """Build context for a task based on its retain_full_context setting"""
+        if not (current_task.previous_tasks or current_task.context):
+            return ""
+            
+        context = "\nInput data from previous tasks:"
+        
+        if current_task.retain_full_context:
+            # Original behavior: include all previous tasks
+            for prev_name in current_task.previous_tasks:
+                prev_task = next((t for t in self.tasks.values() if t.name == prev_name), None)
+                if prev_task and prev_task.result:
+                    context += f"\n{prev_name}: {prev_task.result.raw}"
+                    
+            # Add data from context tasks
+            if current_task.context:
+                for ctx_task in current_task.context:
+                    if ctx_task.result and ctx_task.name != current_task.name:
+                        context += f"\n{ctx_task.name}: {ctx_task.result.raw}"
+        else:
+            # New behavior: only include the most recent previous task
+            if current_task.previous_tasks:
+                # Get the most recent previous task (last in the list)
+                prev_name = current_task.previous_tasks[-1]
+                prev_task = next((t for t in self.tasks.values() if t.name == prev_name), None)
+                if prev_task and prev_task.result:
+                    context += f"\n{prev_name}: {prev_task.result.raw}"
+                    
+            # For context tasks, still include the most recent one
+            if current_task.context:
+                # Get the most recent context task with a result
+                for ctx_task in reversed(current_task.context):
+                    if ctx_task.result and ctx_task.name != current_task.name:
+                        context += f"\n{ctx_task.name}: {ctx_task.result.raw}"
+                        break  # Only include the most recent one
+                        
+        return context
+
     def _find_next_not_started_task(self) -> Optional[Task]:
         """Fallback mechanism to find the next 'not started' task."""
         fallback_attempts = 0
@@ -147,25 +185,8 @@ async def aworkflow(self) -> AsyncGenerator[str, None]:
             """)
 
             # Add context from previous tasks to description
-            if current_task.previous_tasks or current_task.context:
-                context = "\nInput data from previous tasks:"
-
-                # Add data from previous tasks in workflow
-                for prev_name in current_task.previous_tasks:
-                    prev_task = next((t for t in self.tasks.values() if t.name == prev_name), None)
-                    if prev_task and prev_task.result:
-                        # Handle loop data
-                        if current_task.task_type == "loop":
-                            context += f"\n{prev_name}: {prev_task.result.raw}"
-                        else:
-                            context += f"\n{prev_name}: {prev_task.result.raw}"
-
-                # Add data from context tasks
-                if current_task.context:
-                    for ctx_task in current_task.context:
-                        if ctx_task.result and ctx_task.name != current_task.name:
-                            context += f"\n{ctx_task.name}: {ctx_task.result.raw}"
-
+            context = self._build_task_context(current_task)
+            if context:
                 # Update task description with context
                 current_task.description = current_task.description + context
 
@@ -778,25 +799,8 @@ def workflow(self):
             """)
 
             # Add context from previous tasks to description
-            if current_task.previous_tasks or current_task.context:
-                context = "\nInput data from previous tasks:"
-
-                # Add data from previous tasks in workflow
-                for prev_name in current_task.previous_tasks:
-                    prev_task = next((t for t in self.tasks.values() if t.name == prev_name), None)
-                    if prev_task and prev_task.result:
-                        # Handle loop data
-                        if current_task.task_type == "loop":
-                            context += f"\n{prev_name}: {prev_task.result.raw}"
-                        else:
-                            context += f"\n{prev_name}: {prev_task.result.raw}"
-
-                # Add data from context tasks
-                if current_task.context:
-                    for ctx_task in current_task.context:
-                        if ctx_task.result and ctx_task.name != current_task.name:
-                            context += f"\n{ctx_task.name}: {ctx_task.result.raw}"
-
+            context = self._build_task_context(current_task)
+            if context:
                 # Update task description with context
                 current_task.description = current_task.description + context
 

src/praisonai-agents/praisonaiagents/task/task.py

@@ -39,7 +39,8 @@ def __init__(
         memory=None,
         quality_check=True,
         input_file: Optional[str] = None,
-        rerun: bool = False # Renamed from can_rerun and logic inverted, default True for backward compatibility
+        rerun: bool = False, # Renamed from can_rerun and logic inverted, default True for backward compatibility
+        retain_full_context: bool = False # By default, only use previous task output, not all previous tasks
     ):
         # Add check if memory config is provided
         if memory is not None or (config and config.get('memory_config')):
@@ -78,6 +79,7 @@ def __init__(
         self.memory = memory
         self.quality_check = quality_check
         self.rerun = rerun # Assigning the rerun parameter
+        self.retain_full_context = retain_full_context
 
         # Set logger level based on config verbose level
         verbose = self.config.get("verbose", 0)

tests/unit/test_context_management.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""
+Test script to verify context management functionality
+"""
+
+import sys
+import os
+sys.path.append(os.path.join(os.path.dirname(__file__), 'src', 'praisonai-agents'))
+
+from praisonaiagents.task.task import Task
+from praisonaiagents.agent.agent import Agent
+from praisonaiagents.process.process import Process
+from praisonaiagents.main import TaskOutput
+
+def test_context_management():
+    """Test context management with and without retain_full_context"""
+    
+    print("Starting context management test...")
+    
+    try:
+        # Create a mock agent
+        agent = Agent(name="Test Agent", role="Tester", goal="Test context management")
+        print("✓ Created test agent")
+        
+        # Create tasks with results
+        task1 = Task(
+            name="task1",
+            description="First task",
+            agent=agent,
+            status="completed"
+        )
+        task1.result = TaskOutput(
+            description="First task",
+            raw="Result from task 1",
+            agent="Test Agent"
+        )
+        print("✓ Created task1")
+        
+        task2 = Task(
+            name="task2", 
+            description="Second task",
+            agent=agent,
+            status="completed"
+        )
+        task2.result = TaskOutput(
+            description="Second task",
+            raw="Result from task 2",
+            agent="Test Agent"
+        )
+        # Set up the previous_tasks manually since it's not a constructor parameter
+        task2.previous_tasks = ["task1"]
+        print("✓ Created task2")
+        
+        # Test case 1: Default behavior (retain_full_context=False)
+        task3_limited = Task(
+            name="task3_limited",
+            description="Third task with limited context",
+            agent=agent,
+            retain_full_context=False  # Default behavior
+        )
+        # Set up the previous_tasks manually
+        task3_limited.previous_tasks = ["task1", "task2"]
+        print("✓ Created task3_limited")
+        
+        # Test case 2: Full context retention (retain_full_context=True)
+        task3_full = Task(
+            name="task3_full",
+            description="Third task with full context",
+            agent=agent,
+            retain_full_context=True  # Original behavior
+        )
+        # Set up the previous_tasks manually
+        task3_full.previous_tasks = ["task1", "task2"]
+        print("✓ Created task3_full")
+        
+        # Create process and test context building
+        tasks_dict = {
+            task1.id: task1,
+            task2.id: task2,
+            task3_limited.id: task3_limited,
+            task3_full.id: task3_full
+        }
+        
+        process = Process(tasks=tasks_dict, agents=[agent])
+        print("✓ Created process")
+        
+        # Test limited context
+        limited_context = process._build_task_context(task3_limited)
+        print("Limited context (retain_full_context=False):")
+        print(f"'{limited_context}'")
+        print()
+        
+        # Test full context
+        full_context = process._build_task_context(task3_full)
+        print("Full context (retain_full_context=True):")
+        print(f"'{full_context}'")
+        print()
+        
+        # Verify results
+        assert "task2" in limited_context, "Limited context should include most recent task"
+        assert "task1" not in limited_context, "Limited context should NOT include earlier tasks"
+        
+        assert "task1" in full_context, "Full context should include all previous tasks"
+        assert "task2" in full_context, "Full context should include all previous tasks"
+        
+        print("✅ All tests passed!")
+        print("- Limited context only includes most recent previous task")
+        print("- Full context includes all previous tasks")
+        print("- Backwards compatibility maintained")
+        
+    except Exception as e:
+        print(f"❌ Test failed with error: {e}")
+        import traceback
+        traceback.print_exc()
+        raise
+
+if __name__ == "__main__":
+    test_context_management()