Update api_routes_v2.py

api_routes_v2.py

@@ -1023,150 +1023,149 @@ async def get_session_history(
 ):
     """
     Get conversation history for a session.
-    ENHANCEMENT: For pipeline completion messages, add final_output from S3.
+    Correct behavior:
+    - Preserve pipeline_id on messages
+    - Enrich ONLY using stored pipeline_id
+    - No inference, no guessing, no context scanning
     """
     try:
         # Load conversation history from S3
         history = _load_conversation_from_s3(session_id)
-        
-        # NEW: Add final_output to pipeline completion messages
+
         enhanced_history = []
+
         for msg in history:
             msg_copy = msg.copy()
-            
-            # Check if this is a pipeline completion message
-            # Look for the specific pattern in your example
+
+            # ✅ ALWAYS pass through pipeline_id if present
+            pipeline_id = msg.get("pipeline_id")
+            if pipeline_id:
+                msg_copy["pipeline_id"] = pipeline_id
+
+            role = msg_copy.get("role")
             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)
+
+            # ✅ Enrich ONLY pipeline completion messages WITH pipeline_id
+            if (
+                role == "assistant"
+                and "Pipeline completed" in content
+                and pipeline_id
+            ):
+                try:
+                    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")
+                    )
+
+                    # Attach final output if present
+                    if "result" in pipeline_data:
+                        msg_copy["final_output"] = pipeline_data["result"]
+
+                    # Attach 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:
+                    # Non-fatal enrichment failure
+                    msg_copy["final_output_error"] = str(e)
+
+            # Ensure message_id exists
             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):
+
+            # Normalize timestamp
+            if "timestamp" in msg_copy and isinstance(
+                msg_copy["timestamp"], datetime
+            ):
                 msg_copy["timestamp"] = msg_copy["timestamp"].isoformat()
-            
+
             enhanced_history.append(msg_copy)
-        
-        # Apply limit (existing code)
+
+        # Apply limit (keep most recent)
         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 (existing code, but add results)
-        pipelines_hist = s.get("pipelines_history", [])
+
+        # Load session metadata
+        session = session_manager.get_session(session_id) or {}
+
+        # ----- PIPELINES HISTORY (unchanged but cleaned) -----
+        pipelines_hist = session.get("pipelines_history", [])
         enhanced_pipelines = []
+
         for pipeline_meta in pipelines_hist:
             enhanced_pipe = pipeline_meta.copy()
-            
-            # Rename result_preview to result
-            if "result_preview" in enhanced_pipe:
-                enhanced_pipe["result"] = enhanced_pipe.pop("result_preview")
-            
-            # Remove internal S3 keys from response
+
+            # Remove internal-only keys
             enhanced_pipe.pop("pipeline_s3_key", None)
-            
-            # Try to load full results from S3 for pipelines_history too
+
+            # Load full pipeline definition if available
             pipeline_s3_key = pipeline_meta.get("pipeline_s3_key")
             if pipeline_s3_key:
                 try:
-                    resp = s3.get_object(Bucket=S3_BUCKET, Key=pipeline_s3_key)
-                    pipeline_def = json.loads(resp["Body"].read().decode("utf-8"))
-                    
-                    # Extract tools/components
-                    components = pipeline_def.get("components") or pipeline_def.get("pipeline_steps", [])
-                    tools_list = [
-                        comp.get("tool_name", comp.get("tool", "unknown"))
-                        for comp in components
-                    ]
-                    
-                    enhanced_pipe["tools"] = tools_list
-                    enhanced_pipe["component_count"] = len(components)
+                    resp = s3.get_object(
+                        Bucket=S3_BUCKET,
+                        Key=pipeline_s3_key
+                    )
+                    pipeline_def = json.loads(
+                        resp["Body"].read().decode("utf-8")
+                    )
+
+                    components = (
+                        pipeline_def.get("components")
+                        or pipeline_def.get("pipeline_steps", [])
+                    )
+
                     enhanced_pipe["components"] = components
-                    
-                    # Load results if available
+                    enhanced_pipe["component_count"] = len(components)
+                    enhanced_pipe["tools"] = [
+                        c.get("tool_name", c.get("tool", "unknown"))
+                        for c in components
+                    ]
+
                     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"] = []
-                    enhanced_pipe["component_count"] = 0
                     enhanced_pipe["components"] = []
+                    enhanced_pipe["component_count"] = 0
+                    enhanced_pipe["tools"] = []
+                    enhanced_pipe["load_error"] = str(e)
+
             else:
-                enhanced_pipe["tools"] = []
-                enhanced_pipe["component_count"] = 0
                 enhanced_pipe["components"] = []
-            
-            # Add hasError field if not present
+                enhanced_pipe["component_count"] = 0
+                enhanced_pipe["tools"] = []
+
             if "hasError" not in enhanced_pipe:
-                enhanced_pipe["hasError"] = enhanced_pipe.get("status") == "failed"
-            
+                enhanced_pipe["hasError"] = (
+                    enhanced_pipe.get("status") == "failed"
+                )
+
             enhanced_pipelines.append(enhanced_pipe)
-        
-        # Ensure sorting by most recent first
+
+        # Sort pipelines by recency
         enhanced_pipelines.sort(
-            key=lambda p: p.get("updated_at") or p.get("created_at") or "", 
+            key=lambda p: p.get("updated_at") or p.get("created_at") or "",
             reverse=True
         )
 
         return {
             "session_id": session_id,
-            "history": enhanced_history,  # Now includes final_output for pipeline messages
+            "history": enhanced_history,
             "count": len(enhanced_history),
             "limit": limit,
-            "chat_name": s.get("chat_name"),
+            "chat_name": session.get("chat_name"),
             "pipelines_history": enhanced_pipelines
         }
 
@@ -1176,6 +1175,7 @@ async def get_session_history(
             detail=f"Error retrieving session history: {str(e)}"
         )
 
+
 @router.delete("/sessions/{session_id}")
 async def delete_session(session_id: str):
     """