Create query_engine.py

indexes/query_engine.py

@@ -0,0 +1,106 @@
+def query(self, query_text: str) -> Dict[str, Any]:
+    """Process a natural language query across CSV files."""
+    # Find relevant CSV files
+    relevant_csvs = self.index_manager.find_relevant_csvs(query_text)
+    
+    if not relevant_csvs:
+        return {
+            "answer": "No relevant CSV files found for your query.",
+            "sources": []
+        }
+    
+    # Check for direct statistical queries
+    direct_answer = self._handle_statistical_query(query_text, relevant_csvs)
+    if direct_answer:
+        return {
+            "answer": direct_answer,
+            "sources": self._get_sources(relevant_csvs)
+        }
+    
+    # If not a direct statistical query, use the regular approach
+    context = self._prepare_context(query_text, relevant_csvs)
+    prompt = self._generate_prompt(query_text, context)
+    response = self.llm.complete(prompt)
+    
+    return {
+        "answer": response.text,
+        "sources": self._get_sources(relevant_csvs)
+    }
+
+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()
+    
+    # Detect query type
+    is_avg_query = "average" in query_lower or "mean" in query_lower or "avg" in query_lower
+    is_max_query = "maximum" in query_lower or "max" in query_lower
+    is_min_query = "minimum" in query_lower or "min" in query_lower
+    is_count_query = "count" in query_lower or "how many" in query_lower
+    
+    if not (is_avg_query or is_max_query or is_min_query or is_count_query):
+        return None  # Not a statistical query
+    
+    # Extract potential column names from query
+    query_words = set(query_lower.replace("?", "").replace(",", "").split())
+    
+    for csv_id in csv_ids:
+        if csv_id not in self.index_manager.indexes:
+            continue
+            
+        file_path = self.index_manager.indexes[csv_id]["path"]
+        metadata = self.index_manager.indexes[csv_id]["metadata"]
+        
+        try:
+            df = pd.read_csv(file_path)
+            
+            # Find relevant columns based on query
+            target_columns = []
+            for col in df.columns:
+                col_lower = col.lower()
+                # Check if column name appears in query
+                if any(word in col_lower for word in query_words):
+                    target_columns.append(col)
+            
+            # If no direct matches, try to infer from common column names
+            if not target_columns:
+                if "age" in query_lower:
+                    age_cols = [col for col in df.columns if "age" in col.lower()]
+                    if age_cols:
+                        target_columns = age_cols
+                elif "income" in query_lower or "salary" in query_lower:
+                    income_cols = [col for col in df.columns if any(term in col.lower() 
+                                   for term in ["income", "salary", "wage", "earnings"])]
+                    if income_cols:
+                        target_columns = income_cols
+                # Add more common column inferences as needed
+            
+            # If still no matches, use all numeric columns
+            if not target_columns:
+                target_columns = df.select_dtypes(include=['number']).columns.tolist()
+            
+            # Perform the requested calculation
+            results = []
+            for col in target_columns:
+                if not pd.api.types.is_numeric_dtype(df[col]):
+                    continue
+                    
+                if is_avg_query:
+                    value = df[col].mean()
+                    results.append(f"The average {col} is {value:.2f}")
+                elif is_max_query:
+                    value = df[col].max()
+                    results.append(f"The maximum {col} is {value}")
+                elif is_min_query:
+                    value = df[col].min()
+                    results.append(f"The minimum {col} is {value}")
+                elif is_count_query:
+                    value = len(df)
+                    results.append(f"The total count of {col} is {value}")
+            
+            if results:
+                return "\n".join(results)
+                
+        except Exception as e:
+            print(f"Error processing CSV for statistical query: {e}")
+    
+    return None  # No results found