Update insight_and_tasks/agents/task_extraction_agent.py

insight_and_tasks/agents/task_extraction_agent.py

@@ -195,43 +195,56 @@ class TaskExtractionAgent:
 
         try:
             logger.info(f"Running TaskExtractionAgent for user_id: {user_id}, session_id: {session.id}")
-            async for event in runner.run(
+            
+            # Fix: Use regular for loop instead of async for, since runner.run() returns a generator
+            run_result = runner.run(
                 user_id=user_id,
                 session_id=session.id,
                 new_message=user_input_content
-            ):
-                # LlmAgent with output_schema stores the result in event.actions.state_delta[output_key]
-                if (hasattr(event, 'actions') and event.actions and
-                    hasattr(event.actions, 'state_delta') and
-                    isinstance(event.actions.state_delta, dict) and
-                    self.agent.output_key in event.actions.state_delta):
+            )
+            
+            # Check if it's an async iterator or regular generator
+            if hasattr(run_result, '__aiter__'):
+                # It's an async iterator, use async for
+                async for event in run_result:
+                    if (hasattr(event, 'actions') and event.actions and
+                        hasattr(event.actions, 'state_delta') and
+                        isinstance(event.actions.state_delta, dict) and
+                        self.agent.output_key in event.actions.state_delta):
+                        
+                        extracted_data_dict = event.actions.state_delta[self.agent.output_key]
+                        logger.info(f"Successfully extracted structured data via LlmAgent state_delta.")
+                        break
+                    
+                    # Capture text parts for debugging if direct structured output isn't found first
+                    if hasattr(event, 'content') and event.content and event.content.parts:
+                        for part in event.content.parts:
+                            if hasattr(part, 'text'):
+                                 full_response_text_for_debug += part.text
+            else:
+                # It's a regular generator, use regular for loop
+                for event in run_result:
+                    if (hasattr(event, 'actions') and event.actions and
+                        hasattr(event.actions, 'state_delta') and
+                        isinstance(event.actions.state_delta, dict) and
+                        self.agent.output_key in event.actions.state_delta):
+                        
+                        extracted_data_dict = event.actions.state_delta[self.agent.output_key]
+                        logger.info(f"Successfully extracted structured data via LlmAgent state_delta.")
+                        break
                     
-                    extracted_data_dict = event.actions.state_delta[self.agent.output_key]
-                    logger.info(f"Successfully extracted structured data via LlmAgent state_delta.")
-                    break # Assuming full structured output comes in one event with state_delta
-                
-                # Capture text parts for debugging if direct structured output isn't found first
-                if hasattr(event, 'content') and event.content and event.content.parts:
-                    for part in event.content.parts:
-                        if hasattr(part, 'text'):
-                             full_response_text_for_debug += part.text
+                    # Capture text parts for debugging if direct structured output isn't found first
+                    if hasattr(event, 'content') and event.content and event.content.parts:
+                        for part in event.content.parts:
+                            if hasattr(part, 'text'):
+                                 full_response_text_for_debug += part.text
 
             if not extracted_data_dict and full_response_text_for_debug:
                 logger.warning("LlmAgent did not produce structured output in state_delta. Raw text response was: %s", 
-                               full_response_text_for_debug[:500] + "...") # Log snippet
-                # Attempt to parse the raw text if it looks like JSON (fallback, not ideal)
-                # This is a basic fallback; robust JSON cleaning might be needed if LLM doesn't adhere to schema.
-                # For now, we rely on LlmAgent's output_schema handling.
+                               full_response_text_for_debug[:500] + "...") 
 
         except Exception as e:
             logger.error(f"Error during TaskExtractionAgent execution: {e}", exc_info=True)
-            # Fallback to returning an empty TaskExtractionOutput with error info
-            return TaskExtractionOutput(
-                current_quarter_info=f"Q{self._get_quarter(self.current_date)}, {self._days_until_quarter_end(self.current_date)} days remaining",
-                okrs=[],
-                overall_strategic_focus=f"Error during task extraction: {e}",
-                generation_timestamp=datetime.utcnow().isoformat()
-            )
         finally:
             try:
                 await runner.session_service.delete_session(