Update indexes/query_engine.py

indexes/query_engine.py

@@ -1,106 +1,118 @@
-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:
+from typing import Dict, List, Any, Optional
+import pandas as pd
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+class CSVQueryEngine:
+
+    def __init__(self, index_manager, llm):
+        """Initialize with index manager and language model."""
+        self.index_manager = index_manager
+        self.llm = llm
+
+    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": direct_answer,
+            "answer": response.text,
             "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"]
+    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()
         
-        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}")
+        # 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"]
             
-            if results:
-                return "\n".join(results)
+            try:
+                df = pd.read_csv(file_path)
                 
-        except Exception as e:
-            print(f"Error processing CSV for statistical query: {e}")
-    
-    return None  # No results found
+                # 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