Update data_analysis_agent.py

data_analysis_agent.py

@@ -612,7 +612,7 @@ Use types: histogram, bar, scatter, heatmap, line"""
             insights = state["insights"]
             dataset_info = state["dataset_info"]
             
-            # MUCH SIMPLER prompt - just ask for a numbered list
+            # Simple prompt that works well
             prompt = f"""Generate 5 recommendations for this dataset:
 
 Dataset: {dataset_info.get('shape', [0])[0]} rows, {dataset_info.get('shape', [0])[1]} columns
@@ -635,21 +635,58 @@ Write exactly 5 numbered recommendations:
             logger.info(response_content)
             logger.info("=" * 50)
             
-            # FORCE create 5 recommendations regardless of what Groq returns
+            # PROPER PARSING: Extract the 5 numbered recommendations
             recommendations = []
             
-            # Try to parse what we got
+            # Split by lines and look for numbered items
             lines = response_content.split('\n')
+            current_rec = ""
+            current_num = 0
+            
             for line in lines:
                 line = line.strip()
-                if line and len(line) > 20:  # Any substantial line
-                    if not line.startswith('**Recommendation'):
-                        line = f"**Recommendation {len(recommendations) + 1}:** {line}"
-                    recommendations.append(line)
-                    if len(recommendations) >= 5:
-                        break
-            
-            # If we still don't have 5, force create them
+                
+                # Check if this line starts a new numbered recommendation
+                if line.startswith(('1.', '2.', '3.', '4.', '5.')):
+                    # Save the previous recommendation if we have one
+                    if current_rec and current_num > 0:
+                        clean_rec = current_rec.strip()
+                        if len(clean_rec) > 20:
+                            recommendations.append(f"**Recommendation {current_num}:** {clean_rec}")
+                    
+                    # Start new recommendation
+                    current_num = int(line[0])  # Get the number
+                    current_rec = line[2:].strip()  # Get text after "1. "
+                    
+                elif current_rec and line:  # Continue previous recommendation
+                    current_rec += " " + line
+            
+            # Don't forget the last recommendation
+            if current_rec and current_num > 0:
+                clean_rec = current_rec.strip()
+                if len(clean_rec) > 20:
+                    recommendations.append(f"**Recommendation {current_num}:** {clean_rec}")
+            
+            logger.info(f"PARSED RECOMMENDATIONS COUNT: {len(recommendations)}")
+            
+            # If parsing didn't work well, try a simpler approach
+            if len(recommendations) < 3:
+                logger.warning("Primary parsing failed, trying alternative...")
+                
+                # Alternative: Just look for lines that start with numbers
+                recommendations = []
+                for line in lines:
+                    line = line.strip()
+                    if line and line[0].isdigit() and '. ' in line[:5]:
+                        # Extract the text after the number
+                        rec_text = line.split('. ', 1)[1] if '. ' in line else line[2:]
+                        if len(rec_text) > 20:
+                            rec_num = len(recommendations) + 1
+                            recommendations.append(f"**Recommendation {rec_num}:** {rec_text}")
+                        if len(recommendations) >= 5:
+                            break
+            
+            # Final fallback - ensure we have 5 recommendations
             while len(recommendations) < 5:
                 rec_num = len(recommendations) + 1
                 fallback_recs = [
@@ -665,13 +702,13 @@ Write exactly 5 numbered recommendations:
                 else:
                     recommendations.append(f"**Recommendation {rec_num}:** Conduct additional analysis to identify optimization opportunities")
             
-            # FORCE exactly 5 recommendations
+            # Ensure exactly 5 recommendations
             recommendations = recommendations[:5]
             
             # LOG what we're returning
             logger.info(f"FINAL RECOMMENDATIONS COUNT: {len(recommendations)}")
             for i, rec in enumerate(recommendations, 1):
-                logger.info(f"REC {i}: {rec}")
+                logger.info(f"FINAL REC {i}: {rec}")
             
             state["recommendations"] = recommendations
             state["current_step"] = "recommendation_engine"