Update api_routes_v2.py

api_routes_v2.py

@@ -749,7 +749,8 @@ def _assistant_response_payload(
     output: Optional[Dict[str, Any]] = None,
     final_output: Optional[Dict[str, Any]] = None,
     exception: Optional[str] = None,
-    pipeline_result: Optional[Dict[str, Any]] = None  # NEW parameter
+    pipeline_result: Optional[Dict[str, Any]] = None,
+    pipeline_id: Optional[str] = None,  # ✅ ADD THIS PARAMETER
 ) -> ChatResponse:
     """
     Create ChatResponse payload with all required fields.
@@ -758,12 +759,13 @@ def _assistant_response_payload(
     from services.schemas import generate_message_id
     message_id = generate_message_id()
     
-    # Persist assistant message to S3 WITH pipeline results
+    # Persist assistant message to S3 WITH pipeline results AND pipeline_id
     _add_and_mirror_message(
         chat_id=chat_id,
         role="assistant",
         content=friendly_response,
-        pipeline_result=pipeline_result  # NEW: Pass pipeline results
+        pipeline_result=pipeline_result,  # Pass pipeline results
+        pipeline_id=pipeline_id  # ✅ PASS THE PIPELINE_ID HERE
     )
     
     # Get file metadata from session
@@ -1450,16 +1452,6 @@ async def chat_unified(
                 if is_success:
                     friendly = "🎉 Pipeline completed successfully!"
                     
-                    # ✅ NEW: persist assistant message with pipeline result
-                    _add_and_mirror_message(
-                        chat_id,
-                        "assistant",
-                        friendly,
-                        pipeline_id=pipeline_id,
-                        pipeline_status="completed",
-                        pipeline_result=result
-                    )
-                    
                     output = {
                         "component_summary": "Pipeline executed successfully",
                         "steps": total_steps,
@@ -1514,7 +1506,8 @@ async def chat_unified(
                         output=output,
                         final_output=final_output,
                         exception=exception_msg,
-                        pipeline_result=result  # NEW: Pass the pipeline results
+                        pipeline_result=result,  # NEW: Pass the pipeline results
+                        pipeline_id=pipeline_id  # ✅ PASS pipeline_id HERE
                     )
                 else:
                     # Pipeline failed or partially completed
@@ -1536,16 +1529,6 @@ async def chat_unified(
                     else:
                         friendly = f"⚠️ Pipeline partially completed: {error_msg}"
                     
-                    # ✅ NEW: persist failure result
-                    _add_and_mirror_message(
-                        chat_id,
-                        "assistant",
-                        friendly,
-                        pipeline_id=pipeline_id,
-                        pipeline_status=pipeline_status,
-                        pipeline_result=result
-                    )
-                    
                     output = {
                         "component_summary": f"Pipeline {pipeline_status}",
                         "steps": total_steps,
@@ -1573,7 +1556,8 @@ async def chat_unified(
                         output=output,
                         final_output=final_output,
                         exception=exception_msg,
-                        pipeline_result=result  # NEW: Pass results even for failures
+                        pipeline_result=result,  # NEW: Pass results even for failures
+                        pipeline_id=pipeline_id  # ✅ PASS pipeline_id HERE
                     )
             except Exception as e:
                 session_manager.update_session(chat_id, {"state": "initial"})
@@ -1609,16 +1593,6 @@ async def chat_unified(
                     "error": str(e)
                 }
                 
-                # ✅ NEW: persist exception result
-                _add_and_mirror_message(
-                    chat_id,
-                    "assistant",
-                    friendly,
-                    pipeline_id=pipeline_id,
-                    pipeline_status="failed",
-                    pipeline_result=error_result
-                )
-                
                 api_data = {
                     "type": "error",
                     "error_code": "PIPELINE_EXECUTION_FAILED",
@@ -1641,7 +1615,8 @@ async def chat_unified(
                     api_data=api_data,
                     state="initial",
                     exception=str(e),
-                    pipeline_result=error_result  # NEW: Include error in pipeline_result
+                    pipeline_result=error_result,  # NEW: Include error in pipeline_result
+                    pipeline_id=pipeline_id  # ✅ PASS pipeline_id HERE
                 )
             finally:
                 try: