Update api_routes_v2.py

api_routes_v2.py

@@ -948,50 +948,102 @@ async def get_session_history(
     limit: int = 50
 ):
     """
-    Get conversation history for a session
-    CHANGE: Adds 'chat_name' and 'pipelines_history' (added fields only).
+    Get conversation history for a session.
+    ENHANCEMENT: For pipeline completion messages, add final_output from S3.
     """
     try:
-        # FIXED: Load conversation history from S3 (where messages are actually stored)
-        # Previously was using session_manager.get_session_history() which reads from 
-        # MongoDB messages collection, but messages are stored in S3 conversation file
+        # Load conversation history from S3
         history = _load_conversation_from_s3(session_id)
         
-        # Apply limit
-        if limit and len(history) > limit:
-            history = history[-limit:]  # Get last N messages
-
-        # Format datetime objects to ISO strings for JSON serialization
-        # Also add message_id to each message if not present
-        from services.schemas import generate_message_id
-        formatted_history = []
+        # NEW: Add final_output to pipeline completion messages
+        enhanced_history = []
         for msg in history:
             msg_copy = msg.copy()
             
-            # Add message_id if not present
+            # Check if this is a pipeline completion message
+            # Look for the specific pattern in your example
+            content = msg_copy.get("content", "")
+            role = msg_copy.get("role", "")
+            
+            # Identify pipeline completion messages
+            # Your example: "🎉 Pipeline completed successfully!"
+            if role == "assistant" and "Pipeline completed" in content:
+                # Try to find pipeline_id for this message
+                pipeline_id = None
+                
+                # Method 1: Check message metadata (if we stored it)
+                if "pipeline_id" in msg_copy:
+                    pipeline_id = msg_copy.get("pipeline_id")
+                
+                # Method 2: Look in recent context
+                if not pipeline_id:
+                    # Check the last few assistant messages for pipeline context
+                    for prev_msg in reversed(enhanced_history):
+                        prev_content = prev_msg.get("content", "")
+                        if "Pipeline Created" in prev_content and "output" in prev_msg:
+                            pipeline_id = prev_msg["output"].get("pipeline_id")
+                            break
+                
+                # Load results from S3 if we found a pipeline_id
+                if pipeline_id:
+                    try:
+                        # Load pipeline results from S3
+                        s3_key = f"{S3_PREFIX}/pipelines/{pipeline_id}.json"
+                        resp = s3.get_object(Bucket=S3_BUCKET, Key=s3_key)
+                        pipeline_data = json.loads(resp["Body"].read().decode("utf-8"))
+                        
+                        # Check if results exist
+                        if "result" in pipeline_data:
+                            # Extract the final output from results
+                            results = pipeline_data["result"]
+                            
+                            # Add final_output to the message
+                            msg_copy["final_output"] = results
+                            
+                            # Also store pipeline metadata
+                            msg_copy["pipeline_metadata"] = {
+                                "pipeline_id": pipeline_id,
+                                "pipeline_name": pipeline_data.get("pipeline_name", "Unknown"),
+                                "status": pipeline_data.get("status", "unknown"),
+                                "executed_at": pipeline_data.get("updated_at", pipeline_data.get("created_at"))
+                            }
+                    except Exception as e:
+                        print(f"Warning: Could not load pipeline results for {pipeline_id}: {e}")
+                        # Optionally add error info
+                        msg_copy["final_output_error"] = str(e)
+            
+            # Add message_id if not present (existing code)
             if "message_id" not in msg_copy:
+                from services.schemas import generate_message_id
                 msg_copy["message_id"] = generate_message_id()
             
+            # Format timestamp (existing code)
             if "timestamp" in msg_copy and isinstance(msg_copy["timestamp"], datetime):
                 msg_copy["timestamp"] = msg_copy["timestamp"].isoformat()
-            formatted_history.append(msg_copy)
-
+            
+            enhanced_history.append(msg_copy)
+        
+        # Apply limit (existing code)
+        if limit and len(enhanced_history) > limit:
+            enhanced_history = enhanced_history[-limit:]
+        
+        # Get session info (existing code)
         s = session_manager.get_session(session_id) or {}
         
-        # Enhanced pipelines_history with component metadata, sorted by most recent first
+        # Enhanced pipelines_history (existing code, but add results)
         pipelines_hist = s.get("pipelines_history", [])
         enhanced_pipelines = []
         for pipeline_meta in pipelines_hist:
             enhanced_pipe = pipeline_meta.copy()
             
-            # CHANGE: Rename result_preview to result
+            # Rename result_preview to result
             if "result_preview" in enhanced_pipe:
                 enhanced_pipe["result"] = enhanced_pipe.pop("result_preview")
             
-            # CHANGE: Remove internal S3 keys from response
+            # Remove internal S3 keys from response
             enhanced_pipe.pop("pipeline_s3_key", None)
             
-            # Load pipeline definition from S3 to get components/tools
+            # Try to load full results from S3 for pipelines_history too
             pipeline_s3_key = pipeline_meta.get("pipeline_s3_key")
             if pipeline_s3_key:
                 try:
@@ -1007,7 +1059,12 @@ async def get_session_history(
                     
                     enhanced_pipe["tools"] = tools_list
                     enhanced_pipe["component_count"] = len(components)
-                    enhanced_pipe["components"] = components  # Full component details
+                    enhanced_pipe["components"] = components
+                    
+                    # Load results if available
+                    if "result" in pipeline_def:
+                        enhanced_pipe["execution_results"] = pipeline_def["result"]
+                        
                 except Exception as e:
                     print(f"Warning: Failed to load pipeline {pipeline_s3_key}: {e}")
                     enhanced_pipe["tools"] = []
@@ -1018,26 +1075,25 @@ async def get_session_history(
                 enhanced_pipe["component_count"] = 0
                 enhanced_pipe["components"] = []
             
-            # CHANGE: Add hasError field if not present
+            # Add hasError field if not present
             if "hasError" not in enhanced_pipe:
                 enhanced_pipe["hasError"] = enhanced_pipe.get("status") == "failed"
             
             enhanced_pipelines.append(enhanced_pipe)
         
-        # Ensure sorting by most recent first (updated_at or created_at)
+        # Ensure sorting by most recent first
         enhanced_pipelines.sort(
             key=lambda p: p.get("updated_at") or p.get("created_at") or "", 
             reverse=True
         )
 
-        # Return with session_id included
         return {
             "session_id": session_id,
-            "history": formatted_history,
-            "count": len(formatted_history),
+            "history": enhanced_history,  # Now includes final_output for pipeline messages
+            "count": len(enhanced_history),
             "limit": limit,
-            "chat_name": s.get("chat_name"),                    # CHANGE
-            "pipelines_history": enhanced_pipelines             # CHANGE: enhanced with components
+            "chat_name": s.get("chat_name"),
+            "pipelines_history": enhanced_pipelines
         }
 
     except Exception as e:
@@ -1323,7 +1379,8 @@ async def chat_unified(
                     friendly = "🎉 Pipeline completed successfully!"
                     output = {
                         "component_summary": "Pipeline executed successfully",
-                        "steps": total_steps
+                        "steps": total_steps,
+                        "pipeline_id": pipeline_id
                     }
                     api_type = "pipeline_completed"
                     exception_msg = None