Refactor: Use MongoDB _id for messages, pipeline_id_index for components, standardize outputs to result:{text:{}} format

services/pipeline_manager.py

@@ -144,9 +144,10 @@ class PipelineManager:
         pipeline_steps = pipeline_definition.get("components", []) or pipeline_definition.get("pipeline_steps", [])
         components = []
         
+        # Generate component IDs in format: execution_id_index
         for i, step in enumerate(pipeline_steps):
             component = ComponentMetadata(
-                component_id=generate_component_id(),
+                component_id=generate_component_id(execution_id, i),
                 component_name=step.get("tool_name", "unknown"),
                 order=i + 1,
                 status=ComponentStatus.PENDING
@@ -193,7 +194,7 @@ class PipelineManager:
             execution_id: Pipeline execution ID
             component_name: Name of component
             status: "executing", "completed", or "failed"
-            output: Component output (stored temporarily, compiled to final output later)
+            output: Component output (will be wrapped in result:{text:{}} format)
             error: Error message if failed
             success_message: Success message if completed
         
@@ -229,9 +230,18 @@ class PipelineManager:
                 else:
                     comp["hasError"] = False
                 
-                # Store component output
+                # Wrap component output in standardized format: result:{text:{}}
                 if output:
-                    comp["component_output"] = output
+                    # Check if already wrapped
+                    if "result" in output and "text" in output.get("result", {}):
+                        comp["component_output"] = output
+                    else:
+                        # Wrap in standardized format
+                        comp["component_output"] = {
+                            "result": {
+                                "text": output
+                            }
+                        }
                 
                 component_found = True
                 break
@@ -252,9 +262,10 @@ class PipelineManager:
         
         # Store output temporarily (will be compiled into final_output later)
         if output:
-            # Store in a temporary location, will be moved to final output
+            # Wrap output before storing
+            wrapped_output = comp["component_output"]  # Already wrapped above
             temp_key = f"sessions/{pipeline['session_id']}/outputs/{execution_id}/temp_{component_name}.json"
-            self.s3.upload_json(temp_key, output, add_prefix=False)
+            self.s3.upload_json(temp_key, wrapped_output, add_prefix=False)
         
         return True
     
@@ -291,32 +302,53 @@ class PipelineManager:
                 workflow_status = "failed"
                 break
         
-        # Get last successful node output
+        # Get last successful node output - extract from result:{text:{}} format
         last_node_output = None
         for result in reversed(components_results):
             if result.get("status") == "completed" and result.get("result"):
-                # Extract text from result
+                # Extract text from standardized result:{text:{}} format
                 res = result.get("result", {})
                 if isinstance(res, dict):
-                    last_node_output = res.get("text") or res.get("summary") or json.dumps(res)[:500]
+                    # Check for result:{text:{}} structure
+                    if "result" in res and "text" in res.get("result", {}):
+                        text_content = res["result"]["text"]
+                        if isinstance(text_content, dict):
+                            last_node_output = text_content.get("text") or text_content.get("summary") or json.dumps(text_content)[:500]
+                        else:
+                            last_node_output = str(text_content)[:500]
+                    else:
+                        # Fallback for non-wrapped format
+                        last_node_output = res.get("text") or res.get("summary") or json.dumps(res)[:500]
                 elif isinstance(res, str):
                     last_node_output = res
                 break
         
-        # Build final output for S3
+        # Build final output for S3 with standardized format
         component_outputs = []
         for i, result in enumerate(components_results):
+            # Ensure result is wrapped in result:{text:{}} format
+            result_data = result.get("result", {})
+            if isinstance(result_data, dict) and "result" in result_data and "text" in result_data.get("result", {}):
+                wrapped_result = result_data
+            else:
+                # Wrap in standardized format
+                wrapped_result = {
+                    "result": {
+                        "text": result_data
+                    }
+                }
+            
             comp_output = ComponentOutputS3(
-                component_id=result.get("component_id", generate_component_id()),
+                component_id=result.get("component_id", f"{execution_id}_{i}"),
                 name=result.get("component_name", f"component_{i+1}"),
                 order=i + 1,
                 status=result.get("status", "unknown"),
-                component_output=result.get("result"),
+                component_output=wrapped_result,
                 hasError=result.get("status") == "failed",
                 error={"message": result.get("error")} if result.get("error") else None,
                 success_message=result.get("success_message"),
                 metadata=result.get("metadata"),
-                result=result.get("result")
+                result=wrapped_result
             )
             component_outputs.append(comp_output.model_dump())
         
@@ -362,7 +394,8 @@ class PipelineManager:
             "final_output_expires_at": presigned["presigned_expires_at"],
             "final_output_s3_key": final_output_key,
             "workflow_status": workflow_status,
-            "last_node_output": last_node_output
+            "last_node_output": last_node_output,
+            "result": {"text": {"summary": last_node_output}} if last_node_output else None
         }
     
     def mark_pipeline_failed(

services/schemas.py

@@ -248,14 +248,9 @@ class PipelineDefinitionS3(BaseModel):
 # UTILITY FUNCTIONS
 # ========================
 
-def generate_message_id() -> str:
-    """Generate unique message ID"""
-    return f"msg_{uuid.uuid4().hex[:12]}"
-
-
-def generate_component_id() -> str:
-    """Generate unique component ID"""
-    return f"comp_{uuid.uuid4().hex[:12]}"
+def generate_component_id(pipeline_id: str, index: int) -> str:
+    """Generate component ID in format: pipeline_id_index"""
+    return f"{pipeline_id}_{index}"
 
 
 def generate_output_id() -> str:

services/session_manager.py

@@ -20,8 +20,7 @@ from pymongo.errors import DuplicateKeyError, ConnectionFailure
 from services.s3_manager import get_s3_manager
 from services.schemas import (
     SessionSchema, MessageSchema, SessionState, MessageRole,
-    SessionStats, MessageS3Content, schema_to_dict, datetime_to_iso,
-    generate_message_id
+    SessionStats, MessageS3Content, schema_to_dict, datetime_to_iso
 )
 
 
@@ -92,13 +91,14 @@ class SessionManager:
         session_id: str,
         role: str,
         content: str,
-        file_data: Optional[Dict[str, Any]] = None
+        file_data: Optional[Dict[str, Any]] = None,
+        result: Optional[Dict[str, Any]] = None
     ) -> Dict[str, Any]:
         """
         Add message to session
         
         Stores:
-        - Metadata in messages collection
+        - Metadata in messages collection (using MongoDB _id as message_id)
         - Full content in S3
         - Updates session last_activity
         
@@ -107,45 +107,62 @@ class SessionManager:
             role: "user", "assistant", or "system"
             content: Message content
             file_data: Optional file metadata
+            result: Optional pipeline result data (for pipeline completion messages)
         
         Returns:
-            Message metadata dict
+            Message metadata dict with MongoDB _id as message_id
         """
-        message_id = generate_message_id()
         now = datetime.utcnow()
         
-        # Store full message content in S3
+        # Create message metadata record WITHOUT message_id (MongoDB will generate _id)
+        content_preview = content[:200] if len(content) > 200 else content
+        
+        message_record = {
+            "session_id": session_id,
+            "role": role,
+            "timestamp": now.isoformat() + "Z",
+            "content_preview": content_preview,
+            "has_file": bool(file_data and file_data.get("file_id")),
+            "file_id": file_data.get("file_id") if file_data else None,
+            "metadata": {}
+        }
+        
+        # Add result if provided (for pipeline completion messages)
+        if result:
+            message_record["result"] = result
+        
+        # Insert message metadata and get the MongoDB _id
+        insert_result = self.messages_collection.insert_one(message_record)
+        message_id = str(insert_result.inserted_id)
+        
+        # Now store full message content in S3 using the MongoDB _id
         s3_key = f"sessions/{session_id}/messages/{message_id}.json"
         s3_bucket = self.s3.bucket_name
         
-        message_s3_content = MessageS3Content(
-            message_id=message_id,
-            role=role,
-            content=content,
-            timestamp=now.isoformat() + "Z",
-            file_data=file_data or {}
-        )
+        message_s3_content = {
+            "message_id": message_id,
+            "role": role,
+            "content": content,
+            "timestamp": now.isoformat() + "Z",
+            "file_data": file_data or {}
+        }
         
-        self.s3.upload_json(s3_key, message_s3_content.model_dump(), add_prefix=False)
+        if result:
+            message_s3_content["result"] = result
         
-        # Create message metadata record
-        content_preview = content[:200] if len(content) > 200 else content
+        self.s3.upload_json(s3_key, message_s3_content, add_prefix=False)
         
-        message_record = MessageSchema(
-            message_id=message_id,
-            session_id=session_id,
-            role=MessageRole(role),
-            timestamp=now,
-            s3_key=s3_key,
-            s3_bucket=s3_bucket,
-            content_preview=content_preview,
-            has_file=bool(file_data and file_data.get("file_id")),
-            file_id=file_data.get("file_id") if file_data else None
+        # Update the message record with S3 info
+        self.messages_collection.update_one(
+            {"_id": insert_result.inserted_id},
+            {
+                "$set": {
+                    "s3_key": s3_key,
+                    "s3_bucket": s3_bucket
+                }
+            }
         )
         
-        # Insert message metadata
-        self.messages_collection.insert_one(schema_to_dict(message_record))
-        
         # Update session last_activity and message count
         self.sessions_collection.update_one(
             {"session_id": session_id},
@@ -155,7 +172,6 @@ class SessionManager:
             }
         )
         
-        
         # Auto-generate chat name after first user message (only if not already generated)
         if role == "user" and not content.lower().startswith("uploaded file"):
             # Check if chat name already exists to avoid unnecessary generation attempts
@@ -166,8 +182,9 @@ class SessionManager:
             if session and not session.get("chat_name"):
                 self._maybe_generate_chat_name(session_id)
         
-        
-        return schema_to_dict(message_record)
+        # Return message record with _id as message_id
+        message_record["message_id"] = message_id
+        return message_record
     
     def get_messages(
         self,
@@ -184,14 +201,17 @@ class SessionManager:
             include_content: Whether to fetch full content from S3
         
         Returns:
-            List of message dicts
+            List of message dicts with _id converted to message_id
         """
-        # Get message metadata from MongoDB
+        # Get message metadata from MongoDB (include _id this time)
         messages = list(self.messages_collection.find(
-            {"session_id": session_id},
-            {"_id": 0}
+            {"session_id": session_id}
         ).sort("timestamp", 1).limit(limit))
         
+        # Convert MongoDB _id to message_id string
+        for msg in messages:
+            msg["message_id"] = str(msg.pop("_id"))
+        
         if not include_content:
             return messages
         
@@ -201,6 +221,9 @@ class SessionManager:
                 full_content = self.s3.download_json(msg["s3_key"], add_prefix=False)
                 msg["content"] = full_content.get("content", msg["content_preview"])
                 msg["file_data"] = full_content.get("file_data", {})
+                # Include result if present in S3
+                if "result" in full_content:
+                    msg["result"] = full_content["result"]
             except Exception as e:
                 # If S3 fetch fails, use preview
                 msg["content"] = msg["content_preview"]