🔧 Fix sample data loader to support window KG files for Replay button

backend/database/sample_data.py

@@ -100,18 +100,61 @@ class SampleDataLoader:
             self._knowledge_graphs = []
             
             for i, sample in enumerate(config["samples"]):
-                # Load the knowledge graph data
+                # Load the main knowledge graph data
                 kg_data = self._load_knowledge_graph(sample["knowledge_graph_file"])
                 
-                # Convert to the expected format
+                # Check if this sample supports replay (has window KGs)
+                supports_replay = sample.get("supports_replay", False)
+                window_info = sample.get("window_info", {})
+                
+                # Convert main KG to the expected format
                 kg_entry = {
                     "filename": sample["knowledge_graph_file"].split("/")[-1],  # Get just the filename
                     "trace_index": i,  # Links to trace by index
                     "graph_data": kg_data["graph_data"]
                 }
+                
+                # Add window metadata for final KG (window_index=None, window_total=count)
+                if supports_replay and window_info:
+                    kg_entry["window_total"] = window_info.get("window_count", 0)
+                    kg_entry["processing_run_id"] = window_info.get("processing_run_id")
+                    logger.debug(f"Main KG {kg_entry['filename']} configured with {kg_entry['window_total']} windows")
+                
                 self._knowledge_graphs.append(kg_entry)
                 
-            logger.info(f"Loaded {len(self._knowledge_graphs)} sample knowledge graphs")
+                # Load window KGs if they exist
+                if supports_replay and window_info.get("window_files"):
+                    window_files = window_info["window_files"]
+                    window_count = window_info.get("window_count", len(window_files))
+                    processing_run_id = window_info.get("processing_run_id")
+                    
+                    logger.debug(f"Loading {len(window_files)} window KGs for {sample['id']}")
+                    
+                    for window_index, window_file in enumerate(window_files):
+                        try:
+                            # Load window KG data
+                            window_kg_data = self._load_knowledge_graph(window_file)
+                            
+                            # Convert window KG to the expected format
+                            window_kg_entry = {
+                                "filename": window_file.split("/")[-1],  # Get just the filename
+                                "trace_index": i,  # Links to same trace
+                                "graph_data": window_kg_data["graph_data"],
+                                "window_index": window_index,  # This makes it a window KG
+                                "window_total": window_count,
+                                "processing_run_id": processing_run_id,
+                                "window_start_char": window_kg_data.get("window_start_char"),
+                                "window_end_char": window_kg_data.get("window_end_char")
+                            }
+                            
+                            self._knowledge_graphs.append(window_kg_entry)
+                            logger.debug(f"Loaded window KG {window_index}: {window_kg_entry['filename']}")
+                            
+                        except Exception as e:
+                            logger.error(f"Failed to load window KG {window_file}: {e}")
+                            continue
+                
+            logger.info(f"Loaded {len(self._knowledge_graphs)} sample knowledge graphs (including window KGs)")
         
         return self._knowledge_graphs
     
@@ -231,17 +274,33 @@ def insert_sample_data(session, force_insert=False):
             try:
                 trace_index = kg_data["trace_index"]
                 if trace_index < len(trace_ids):
+                    # Extract window information from the KG data
+                    window_index = kg_data.get("window_index")  # None for final KG, index for window KG
+                    window_total = kg_data.get("window_total", 1)  # Use provided window_total or default to 1
+                    window_start_char = kg_data.get("window_start_char")
+                    window_end_char = kg_data.get("window_end_char")
+                    processing_run_id = kg_data.get("processing_run_id")
+                    
                     save_knowledge_graph(
                         session=session,
                         filename=kg_data["filename"],
                         graph_data=kg_data["graph_data"],
                         trace_id=trace_ids[trace_index],
-                        window_index=None,  # None for final KG
-                        window_total=1,     # Not None to indicate it's a processed final KG
+                        window_index=window_index,
+                        window_total=window_total,
+                        window_start_char=window_start_char,
+                        window_end_char=window_end_char,
+                        processing_run_id=processing_run_id,
                         is_original=True
                     )
                     results["knowledge_graphs_inserted"] += 1
-                    logger.info(f"Inserted sample knowledge graph: {kg_data['filename']}")
+                    
+                    # Log different messages for final vs window KGs
+                    if window_index is None:
+                        logger.info(f"Inserted sample knowledge graph: {kg_data['filename']} (final, {window_total} windows)")
+                    else:
+                        logger.info(f"Inserted sample knowledge graph: {kg_data['filename']} (window {window_index})")
+                        
             except Exception as e:
                 error_msg = f"Error inserting knowledge graph {kg_data['filename']}: {str(e)}"
                 logger.error(error_msg)