Update app.py

Fixed a bug where the user input instead of the final (approved by the critic) code was sent to the documentation agent which caused it to produce documentation that did not align with the final result.
It should be better now!

app.py

@@ -49,18 +49,29 @@ async def async_stream_task(task_message, api_key, primary_system_message, criti
 
     llm_config = create_llm_config(api_key)
     team, model_client = create_team(llm_config, primary_system_message, critic_system_message)
+
     documentation_triggered = False  # Track if documentation agent was triggered
     final_output = None  # Store the final approved output
 
+    # Variable to store the latest code snippet from the primary agent
+    latest_primary_code = ""
+
     try:
         async for message in team.run_stream(task=task_message):
             if hasattr(message, "source") and hasattr(message, "content"):
+
+                # If it's from the primary agent, store the latest code snippet
+                if message.source == "primary":
+                    latest_primary_code = message.content
+
                 # Handle critic's approval
                 if message.source == "critic" and "APPROVE" in message.content:
                     print("Critic approved the response. Handing off to Documentation Agent...")
                     documentation_triggered = True
-                    final_output = task_message  # Capture the final approved output
+                    # Capture the primary agent's final code
+                    final_output = latest_primary_code
                     break
+
                 yield message.source, message.content
 
         # Trigger Documentation Agent if approved
@@ -124,7 +135,7 @@ iface = gr.Interface(
     description="""Collaborative workflow between Primary, Critic, and Documentation agents.
     1. The user can send a prompt to the primary agent.
     2. The response will then be evaluated by the critic, which either sends feedback back to the primary agent or gives the APPROVAL sign.
-    3. If the APPROVAL sign is given, the documentation agent is asked to write a short documentation for the code (that has been approved by the critic and generated by the priamry agent.
+    3. If the APPROVAL sign is given, the documentation agent is asked to write a short documentation for the code (that has been approved by the critic and generated by the primary agent).
     4. (Note: There is a hard cap of 10 messages for the critic to approve the output of the primary agent. If it fails to do so the workflow is interrupted to prevent long loops)"""
 )