Update indexes/query_engine.py

indexes/query_engine.py

@@ -39,6 +39,22 @@ class CSVQueryEngine:
             "sources": self._get_sources(relevant_csvs)
         }
     
+    def _get_sources(self, csv_ids: List[str]) -> List[Dict[str, str]]:
+        """Get source information for the response."""
+        sources = []
+        
+        for csv_id in csv_ids:
+            if csv_id not in self.index_manager.indexes:
+                continue
+                
+            metadata = self.index_manager.indexes[csv_id]["metadata"]
+            sources.append({
+                "csv": metadata["filename"],
+                "columns": ", ".join(metadata["columns"][:5]) + ("..." if len(metadata["columns"]) > 5 else "")
+            })
+        
+        return sources
+    
     def _handle_statistical_query(self, query: str, csv_ids: List[str]) -> Optional[str]:
         """Handle direct statistical queries without using the LLM."""
         query_lower = query.lower()