ww

app.py

@@ -164,6 +164,111 @@ class DataContextManager:
         except Exception as e:
             logger.error(f"Error in specific dataset search: {e}")
             return []
+        
+        
+    def _search_water_chemicals(self, query: str, df: pd.DataFrame, top_k: int = 5) -> List[Dict[str, Any]]:
+        """Special search method for water_chemicals dataset"""
+        query_lower = query.lower()
+        
+        # Define chemical mappings for better matching
+        chemical_mappings = {
+            'ph': ['pH', 'ph'],
+            'nitrite': ['Nitrite', 'nitrite'],
+            'oxygen': ['Dissolved oxygen', 'oxygen', 'O2'],
+            'dissolved oxygen': ['Dissolved oxygen'],
+            'nitrogen': ['Nitrogen', 'nitrogen', 'N'],
+            'phosphorus': ['Phosphorus', 'phosphorus', 'P'],
+            'temperature': ['Temperature', 'temperature', 'temp']
+        }
+        
+        # Find relevant chemical types based on query
+        relevant_chemicals = set()
+        for key, chemicals in chemical_mappings.items():
+            if key in query_lower:
+                relevant_chemicals.update(chemicals)
+        
+        # If no specific chemicals found, search more broadly
+        if not relevant_chemicals:
+            # Look for any chemical-related terms in the query
+            for chemical in df['observedPropertyDeterminandLabel'].unique():
+                if pd.notna(chemical) and any(term in chemical.lower() for term in query_lower.split()):
+                    relevant_chemicals.add(chemical)
+        
+        # Filter data by relevant chemicals
+        if relevant_chemicals:
+            chemical_filter = df['observedPropertyDeterminandLabel'].isin(relevant_chemicals)
+            filtered_df = df[chemical_filter]
+        else:
+            filtered_df = df
+        
+        # Also filter by country if mentioned in query
+        countries_in_query = []
+        all_countries = df['Country'].unique()
+        for country in all_countries:
+            if pd.notna(country) and country.lower() in query_lower:
+                countries_in_query.append(country)
+        
+        if countries_in_query:
+            country_filter = filtered_df['Country'].isin(countries_in_query)
+            filtered_df = filtered_df[country_filter]
+        
+        # If we have filtered results, return them
+        if len(filtered_df) > 0:
+            # Sort by most recent year and return top results
+            if 'phenomenonTimeReferenceYear' in filtered_df.columns:
+                filtered_df = filtered_df.sort_values('phenomenonTimeReferenceYear', ascending=False)
+            
+            results = []
+            for idx, row in filtered_df.head(top_k).iterrows():
+                result = {
+                    'dataset': 'water_chemicals',
+                    'score': 1.0,  # High score for direct matches
+                    'row_index': int(idx),
+                    'data': row.to_dict(),
+                    'context_type': 'data'
+                }
+                results.append(result)
+            
+            return results
+        
+        # Fallback to TF-IDF search if no direct matches
+        try:
+            search_texts = []
+            for _, row in df.iterrows():
+                text_parts = [
+                    str(row['Country']) if pd.notna(row['Country']) else '',
+                    str(row['observedPropertyDeterminandLabel']) if pd.notna(row['observedPropertyDeterminandLabel']) else '',
+                    str(row['resultMeanValue']) if pd.notna(row['resultMeanValue']) else ''
+                ]
+                search_texts.append(" ".join(text_parts))
+            
+            if search_texts:
+                vectorizer = TfidfVectorizer(stop_words='english', min_df=1, max_features=3000)
+                tfidf_matrix = vectorizer.fit_transform(search_texts)
+                query_vector = vectorizer.transform([query])
+                
+                similarity_scores = cosine_similarity(query_vector, tfidf_matrix)[0]
+                top_indices = similarity_scores.argsort()[-top_k:][::-1]
+                
+                results = []
+                for idx in top_indices:
+                    if similarity_scores[idx] > 0.01:
+                        result = {
+                            'dataset': 'water_chemicals',
+                            'score': float(similarity_scores[idx]),
+                            'row_index': int(idx),
+                            'data': df.iloc[idx].to_dict(),
+                            'context_type': 'data'
+                        }
+                        results.append(result)
+                
+                return results
+        
+        except Exception as e:
+            logger.error(f"Error in water chemicals TF-IDF search: {e}")
+        
+        return []  
+                
 
     def get_data_statistics(self):
         """Get statistics about loaded datasets"""
@@ -179,40 +284,43 @@ class DataContextManager:
     
 def enhanced_determine_query_type(query: str) -> Tuple[str, Dict[str, Any]]:
     """
-    Enhanced query type determination - ADD THIS FUNCTION
+    Enhanced query type determination - FIXED with real column names
     """
     query = query.lower()
     context = {}
     
-    # Data-specific patterns
-    if re.search(r'\b(ph|ph level|acidity|alkalinity)\b', query):
+    # Chemical/pH patterns - search in water_chemicals data
+    if re.search(r'\b(ph|ph level|acidity|alkalinity|nitrite|dissolved oxygen|chemical|chemicals)\b', query):
         return "data", {
             'dataset_preference': 'water_chemicals',
-            'focus_columns': ['pH', 'ph_level', 'acidity', 'alkalinity'],
+            'focus_columns': ['observedPropertyDeterminandLabel', 'resultMeanValue'],
             'data_type': 'chemical'
         }
     
-    elif re.search(r'\b(oxygen|dissolved oxygen|do level|oxygen level)\b', query):
-        return "data", {
-            'dataset_preference': 'water_chemicals', 
-            'focus_columns': ['dissolved_oxygen', 'oxygen_level', 'DO'],
-            'data_type': 'chemical'
-        }
-    
-    elif re.search(r'\b(quality score|water quality|contamination|pollution)\b', query):
+    # Water quality patterns - search in water_quality_scores
+    elif re.search(r'\b(quality score|water quality|overall|rivers|groundwater|open water)\b', query):
         return "data", {
             'dataset_preference': 'water_quality_scores',
-            'focus_columns': ['quality_score', 'contamination_level', 'pollution_index'],
+            'focus_columns': ['Overall', 'Rivers', 'Open water bodies', 'Groundwater'],
             'data_type': 'quality'
         }
     
-    elif re.search(r'\b(inequality|access|drinking water access|water access)\b', query):
+    # Inequality/access patterns - search in drinking_water_inequality  
+    elif re.search(r'\b(inequality|access|drinking water|poorest|richest|ratio)\b', query):
         return "data", {
             'dataset_preference': 'drinking_water_inequality',
-            'focus_columns': ['access_rate', 'inequality_index', 'coverage'],
+            'focus_columns': ['Poorest', 'Richest', 'Ratio'],
             'data_type': 'social'
         }
     
+    # Country-specific queries - search all datasets
+    elif re.search(r'\b(netherlands|usa|america|finland|bulgaria|albania|country|countries)\b', query):
+        return "hybrid", {
+            'search_both': True,
+            'country_query': True
+        }
+    
+    # Statistics/data queries - hybrid search
     elif re.search(r'\b(statistics|stats|numbers|data|metrics|measurements)\b', query):
         return "hybrid", {'search_both': True}
     
@@ -232,7 +340,6 @@ def enhanced_determine_query_type(query: str) -> Tuple[str, Dict[str, Any]]:
     # Default to hybrid search for complex queries
     return "hybrid", {'search_both': True}
 
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,