Update indexes/query_engine.py

indexes/query_engine.py

@@ -38,6 +38,80 @@ class CSVQueryEngine:
         }
     
     def _prepare_context(self, query: str, csv_ids: List[str]) -> str:
+        """Prepare context from relevant CSV files with pre-calculated statistics."""
+        context_parts = []
+        calculated_answers = {}
+        
+        # Check for common statistical questions
+        query_lower = query.lower()
+        is_avg_question = "average" in query_lower or "mean" in query_lower
+        is_max_question = "maximum" in query_lower or "max" in query_lower
+        is_min_question = "minimum" in query_lower or "min" in query_lower
+        
+        # Extract potential column names from query
+        query_words = set(query_lower.replace("?", "").replace(",", "").split())
+        
+        for csv_id in csv_ids:
+            # Get metadata
+            if csv_id not in self.index_manager.indexes:
+                continue
+                
+            metadata = self.index_manager.indexes[csv_id]["metadata"]
+            file_path = self.index_manager.indexes[csv_id]["path"]
+            
+            # Add CSV metadata
+            context_parts.append(f"CSV File: {metadata['filename']}")
+            context_parts.append(f"Columns: {', '.join(metadata['columns'])}")
+            context_parts.append(f"Row Count: {metadata['row_count']}")
+            
+            # Add sample data and calculate statistics
+            try:
+                df = pd.read_csv(file_path)
+                context_parts.append("\nSample Data:")
+                context_parts.append(df.head(3).to_string())
+                
+                # Find relevant columns based on query
+                column_matches = []
+                for col in df.columns:
+                    col_lower = col.lower()
+                    # Check if column name appears in query or is similar to words in query
+                    if col_lower in query_lower or any(word in col_lower for word in query_words):
+                        column_matches.append(col)
+                
+                # If no direct matches, include all numeric columns
+                if not column_matches:
+                    column_matches = df.select_dtypes(include=['number']).columns.tolist()
+                
+                # Calculate statistics for matched columns
+                for col in column_matches:
+                    if pd.api.types.is_numeric_dtype(df[col]):
+                        if is_avg_question:
+                            avg_value = df[col].mean()
+                            context_parts.append(f"\nThe average {col} is: {avg_value:.2f}")
+                            calculated_answers[f"average_{col}"] = avg_value
+                        
+                        if is_max_question:
+                            max_value = df[col].max()
+                            context_parts.append(f"\nThe maximum {col} is: {max_value}")
+                            calculated_answers[f"max_{col}"] = max_value
+                        
+                        if is_min_question:
+                            min_value = df[col].min()
+                            context_parts.append(f"\nThe minimum {col} is: {min_value}")
+                            calculated_answers[f"min_{col}"] = min_value
+                
+            except Exception as e:
+                context_parts.append(f"Error reading CSV: {str(e)}")
+        
+        # Add direct answer if calculated
+        if calculated_answers:
+            context_parts.append("\nDirect Answer:")
+            for key, value in calculated_answers.items():
+                context_parts.append(f"{key.replace('_', ' ')}: {value}")
+        
+        return "\n\n".join(context_parts)
+    
+    def _prepare_context1(self, query: str, csv_ids: List[str]) -> str:
         """Prepare context from relevant CSV files."""
         context_parts = []