Update api_routes_v2.py

api_routes_v2.py

@@ -639,7 +639,8 @@ def _build_api_result_summary(result: Dict[str, Any], session: Dict[str, Any]) -
         metadata["pages_processed"] = pages_processed
     return {"summary": summary, "metadata": metadata}
 
-def _add_and_mirror_message(chat_id: str, role: str, content: str, file_metadata: Optional[Dict[str, Any]] = None):
+def _add_and_mirror_message(chat_id: str, role: str, content: str, file_metadata: Optional[Dict[str, Any]] = None, 
+                           pipeline_result: Optional[Dict[str, Any]] = None):
     """
     V3 RULE: Append message to S3 conversation.
     MongoDB metadata updated via _save_conversation_to_s3.
@@ -648,14 +649,30 @@ def _add_and_mirror_message(chat_id: str, role: str, content: str, file_metadata
     # 1. Load existing
     current_messages = _load_conversation_from_s3(chat_id)
     
-    # 2. Append new
+    # 2. Create message
     new_msg = {
         "role": role,
         "content": content if isinstance(content, str) else json.dumps(content, ensure_ascii=False),
         "timestamp": datetime.utcnow().isoformat() + "Z"
     }
     
-    # V3: Include file metadata if provided
+    # Add result for pipeline completion messages
+    if role == "assistant" and pipeline_result and "pipeline completed" in content.lower():
+        extracted_text = pipeline_result.get("execution_results", {}).get("text", "")
+        if extracted_text:
+            new_msg["result"] = {
+                "text": extracted_text,
+                "pipeline_id": pipeline_result.get("pipeline_id"),
+                "status": pipeline_result.get("status", "completed")
+            }
+        elif pipeline_result.get("error"):
+            new_msg["result"] = {
+                "error": pipeline_result.get("error"),
+                "pipeline_id": pipeline_result.get("pipeline_id"),
+                "status": "failed"
+            }
+    
+    # Add file metadata if provided
     if file_metadata:
         new_msg["file_data"] = {
             "has_file": True,
@@ -682,7 +699,8 @@ def _assistant_response_payload(
     state: str,
     output: Optional[Dict[str, Any]] = None,
     final_output: Optional[Dict[str, Any]] = None,
-    exception: Optional[str] = None
+    exception: Optional[str] = None,
+    pipeline_result: Optional[Dict[str, Any]] = None  # NEW: Add pipeline results
 ) -> ChatResponse:
     """
     Create ChatResponse payload with all required fields.
@@ -691,15 +709,39 @@ def _assistant_response_payload(
     from services.schemas import generate_message_id
     message_id = generate_message_id()
     
-    # Persist assistant message to S3
-    _add_and_mirror_message(chat_id, "assistant", friendly_response)
+    # Create the base message
+    base_message = {
+        "role": "assistant",
+        "content": friendly_response,
+        "timestamp": datetime.utcnow().isoformat() + "Z",
+        "file_data": {"has_file": False}
+    }
+    
+    # NEW: Add pipeline results if available
+    if pipeline_result:
+        # Extract text from pipeline results
+        extracted_text = pipeline_result.get("execution_results", {}).get("text", "")
+        if extracted_text:
+            base_message["result"] = {
+                "text": extracted_text,
+                "pipeline_id": pipeline_result.get("pipeline_id"),
+                "status": pipeline_result.get("status", "completed")
+            }
+    
+    # Persist assistant message to S3 (with results if available)
+    _save_conversation_to_s3(
+        chat_id, 
+        _load_conversation_from_s3(chat_id) + [base_message], 
+        update_activity=True
+    )
     
     # Get file metadata from session
     session = session_manager.get_session(chat_id) or {}
     chat_name = session.get("chat_name")
     file_metadata = session.get("file_metadata", {})
     
-    return ChatResponse(
+    # Create the ChatResponse
+    response = ChatResponse(
         message_id=message_id,
         assistant_response=friendly_response,
         output=output or {},
@@ -715,6 +757,8 @@ def _assistant_response_payload(
         fileName=file_metadata.get("file_name"),
         fileUrl=file_metadata.get("file_url")
     )
+    
+    return response
 
 def parse_s3_uri(uri: str) -> Tuple[str, str]:
     """
@@ -1413,6 +1457,26 @@ async def chat_unified(
                             print(f"Warning: Could not get output_id/download_url: {e}")
                     if not final_output:
                         final_output = None
+                        
+                    # Build api_response data for success
+                    api_data = {
+                        "type": api_type,
+                        "result": api_result,
+                        "pipeline": api_pipeline
+                    }
+                    
+                    # Return response with pipeline result included
+                    return _assistant_response_payload(
+                        chat_id=chat_id,
+                        friendly_response=friendly,
+                        intent={"intent": "pipeline_execute"},
+                        api_data=api_data,
+                        state="initial",
+                        output=output,
+                        final_output=final_output,
+                        exception=exception_msg,
+                        pipeline_result=result  # NEW: Pass the pipeline results
+                    )
                 else:
                     # Pipeline failed or partially completed
                     error_msg = result.get("error", "Pipeline execution incomplete")
@@ -1442,31 +1506,26 @@ async def chat_unified(
                     final_output = {"text": f"Pipeline execution {pipeline_status} with {completed_steps}/{total_steps} steps completed"}
                     api_type = "pipeline_failed" if pipeline_status == "failed" else "pipeline_partial"
                     exception_msg = error_msg
-                
-                # Build api_response data
-                if is_success:
-                    api_data = {
-                        "type": api_type,
-                        "result": api_result,
-                        "pipeline": api_pipeline
-                    }
-                else:
+                    
+                    # Build api_response data for failure
                     api_data = {
                         "type": api_type,
                         "result": result,
                         "pipeline": proposed
                     }
-                
-                return _assistant_response_payload(
-                    chat_id=chat_id,
-                    friendly_response=friendly,
-                    intent={"intent": "pipeline_execute"},
-                    api_data=api_data,
-                    state="initial",
-                    output=output,
-                    final_output=final_output,
-                    exception=exception_msg  # This will set hasError=True when not None
-                )
+                    
+                    # Return response with pipeline result included (even for failures)
+                    return _assistant_response_payload(
+                        chat_id=chat_id,
+                        friendly_response=friendly,
+                        intent={"intent": "pipeline_execute"},
+                        api_data=api_data,
+                        state="initial",
+                        output=output,
+                        final_output=final_output,
+                        exception=exception_msg,
+                        pipeline_result=result  # NEW: Pass results even for failures
+                    )
             except Exception as e:
                 session_manager.update_session(chat_id, {"state": "initial"})
                 pipeline_id = proposed.get("pipeline_id")
@@ -1501,13 +1560,22 @@ async def chat_unified(
                 }
                 if failed_component:
                     api_data["failed_component"] = failed_component
+                
+                # Create error result for the response
+                error_result = {
+                    "error": str(e),
+                    "pipeline_id": pipeline_id,
+                    "status": "failed"
+                }
+                
                 return _assistant_response_payload(
                     chat_id=chat_id,
                     friendly_response=friendly,
                     intent={"intent": "pipeline_execute"},
                     api_data=api_data,
                     state="initial",
-                    exception=str(e)
+                    exception=str(e),
+                    pipeline_result=error_result  # NEW: Include error in pipeline_result
                 )
             finally:
                 try: