Update api_routes_v2.py

api_routes_v2.py

@@ -639,12 +639,16 @@ 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, 
-                           pipeline_result: 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  # NEW: Add pipeline_result parameter
+):
     """
     V3 RULE: Append message to S3 conversation.
-    MongoDB metadata updated via _save_conversation_to_s3.
-    Updates last_activity for all messages (user and assistant interactions).
+    Now includes pipeline results in assistant messages.
     """
     # 1. Load existing
     current_messages = _load_conversation_from_s3(chat_id)
@@ -656,21 +660,63 @@ def _add_and_mirror_message(chat_id: str, role: str, content: str, file_metadata
         "timestamp": datetime.utcnow().isoformat() + "Z"
     }
     
-    # 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"
-            }
+    # NEW: Add pipeline results if available (for assistant messages)
+    if role == "assistant" and pipeline_result:
+        # Check if this is a pipeline completion message
+        if "pipeline completed" in content.lower() or "pipeline failed" in content.lower():
+            # Extract the text/result from pipeline
+            execution_results = pipeline_result.get("execution_results", {})
+            
+            # Try to extract the main result text
+            result_text = ""
+            
+            # Check for image descriptions
+            if execution_results.get("image_descriptions"):
+                image_desc = execution_results["image_descriptions"]
+                if image_desc.get("result") and isinstance(image_desc["result"], list) and len(image_desc["result"]) > 0:
+                    page_result = image_desc["result"][0]
+                    # Use Gemini description if available
+                    if page_result.get("gemini", {}).get("description"):
+                        result_text = page_result["gemini"]["description"]
+                    elif page_result.get("mistral", {}).get("description"):
+                        result_text = page_result["mistral"]["description"]
+            
+            # Check for extracted text (for text extraction pipelines)
+            elif execution_results.get("text"):
+                result_text = execution_results["text"]
+            
+            # Check for component results
+            elif execution_results.get("components_executed"):
+                for comp in execution_results["components_executed"]:
+                    comp_result = comp.get("result", {})
+                    if comp_result.get("text"):
+                        result_text = comp_result["text"]
+                        break
+                    elif comp_result.get("image_descriptions"):
+                        image_desc = comp_result["image_descriptions"]
+                        if image_desc.get("result") and isinstance(image_desc["result"], list) and len(image_desc["result"]) > 0:
+                            page_result = image_desc["result"][0]
+                            if page_result.get("gemini", {}).get("description"):
+                                result_text = page_result["gemini"]["description"]
+                                break
+                            elif page_result.get("mistral", {}).get("description"):
+                                result_text = page_result["mistral"]["description"]
+                                break
+            
+            # Add result to message if we found something
+            if result_text:
+                new_msg["result"] = {
+                    "text": result_text,
+                    "pipeline_id": pipeline_result.get("pipeline_id"),
+                    "status": pipeline_result.get("status", "completed")
+                }
+            elif pipeline_result.get("error"):
+                # For error cases
+                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:
@@ -700,7 +746,7 @@ 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: Add pipeline results
+    pipeline_result: Optional[Dict[str, Any]] = None  # NEW parameter
 ) -> ChatResponse:
     """
     Create ChatResponse payload with all required fields.
@@ -709,30 +755,12 @@ def _assistant_response_payload(
     from services.schemas import generate_message_id
     message_id = generate_message_id()
     
-    # 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
+    # Persist assistant message to S3 WITH pipeline results
+    _add_and_mirror_message(
+        chat_id=chat_id,
+        role="assistant",
+        content=friendly_response,
+        pipeline_result=pipeline_result  # NEW: Pass pipeline results
     )
     
     # Get file metadata from session
@@ -740,8 +768,7 @@ def _assistant_response_payload(
     chat_name = session.get("chat_name")
     file_metadata = session.get("file_metadata", {})
     
-    # Create the ChatResponse
-    response = ChatResponse(
+    return ChatResponse(
         message_id=message_id,
         assistant_response=friendly_response,
         output=output or {},
@@ -757,8 +784,6 @@ 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]:
     """