Optimize MLA lookup and fix tests

backend/gemini_summary.py

@@ -5,15 +5,19 @@
 Includes retry logic with exponential backoff for handling transient failures.
 """
 import os
-import google.generativeai as genai
-from typing import Dict, Optional, Callable, Any
 import warnings
-from async_lru import alru_cache
-from retry_utils import exponential_backoff_retry
-import logging
+from typing import Dict, Optional
+
+# Suppress the Future Warning from google.generativeai
+warnings.filterwarnings("ignore", category=FutureWarning, module="google.generativeai")
+
+import google.generativeai as genai
+
+# Simple in-memory cache to store generated summaries
+# Key: (district, assembly_constituency, mla_name, issue_category)
+# Value: Summary string
+_summary_cache: Dict[tuple, str] = {}
 
-# Configure logging
-logger = logging.getLogger(__name__)
 
 # Suppress deprecation warnings from google.generativeai
 warnings.filterwarnings("ignore", category=FutureWarning, module="google.generativeai")
@@ -97,6 +101,11 @@ async def generate_mla_summary(
     Returns:
         A short paragraph describing the MLA's role and responsibilities
     """
+    # Check cache first
+    cache_key = (district, assembly_constituency, mla_name, issue_category)
+    if cache_key in _summary_cache:
+        return _summary_cache[cache_key]
+
     if not api_key:
         logger.warning("No API key configured, using fallback MLA summary")
         return _get_fallback_summary(mla_name, assembly_constituency, district)
@@ -114,16 +123,13 @@ async def generate_mla_summary(
         Keep it factual, helpful, and encouraging for civic engagement.
         """
 
-        # Generate content without any tools (no Google Search, no internet retrieval)
-        # Explicitly set tools=None to ensure no search/grounding features are used
-        response = await client.aio.models.generate_content(
-            model='gemini-1.5-flash',
-            contents=prompt,
-            config=genai.types.GenerateContentConfig(
-                tools=None  # Explicitly disable all tools including Google Search
-            )
-        )
-        return response.text.strip()
+        response = await model.generate_content_async(prompt)
+        summary = response.text.strip()
+
+        # Store in cache
+        _summary_cache[cache_key] = summary
+
+        return summary
 
     except Exception as e:
         logger.error(f"Gemini Summary Error after all retries: {e}", exc_info=True)

backend/maharashtra_locator.py

@@ -55,7 +55,7 @@ def load_maharashtra_pincode_data() -> Dict[str, Dict[str, Any]]:
     Load and cache Maharashtra pincode to constituency mapping data.
 
     Returns:
-        dict: Dictionary mapping pincode to data
+        dict: Dictionary mapping pincode -> constituency info
     """
     file_path = os.path.join(
         os.path.dirname(__file__),
@@ -65,7 +65,7 @@ def load_maharashtra_pincode_data() -> Dict[str, Dict[str, Any]]:
 
     with open(file_path, "r", encoding="utf-8") as f:
         data_list = json.load(f)
-        # Convert list to dictionary for O(1) lookup
+        # Convert list to dictionary keyed by pincode
         return {item["pincode"]: item for item in data_list}
 
 
@@ -75,7 +75,7 @@ def load_maharashtra_mla_data() -> Dict[str, Dict[str, Any]]:
     Load and cache Maharashtra MLA information data.
 
     Returns:
-        dict: Dictionary mapping constituency to MLA data
+        dict: Dictionary mapping assembly_constituency -> MLA info
     """
     file_path = os.path.join(
         os.path.dirname(__file__),
@@ -85,39 +85,10 @@ def load_maharashtra_mla_data() -> Dict[str, Dict[str, Any]]:
 
     with open(file_path, "r", encoding="utf-8") as f:
         data_list = json.load(f)
-        # Convert list to dictionary for O(1) lookup
+        # Convert list to dictionary keyed by assembly_constituency
         return {item["assembly_constituency"]: item for item in data_list}
 
 
-def get_district_by_pincode_range(pincode: int) -> Optional[str]:
-    """
-    Find district by checking pincode ranges.
-    This is an O(N) fallback where N is number of ranges (~35).
-    """
-    for start, end, district in DISTRICT_RANGES:
-        if start <= pincode <= end:
-            return district
-    return None
-
-
-@lru_cache(maxsize=1)
-def _load_maharashtra_pincode_map() -> Dict[str, Dict[str, Any]]:
-    """
-    Load and cache Maharashtra pincode to constituency mapping as a dict for O(1) lookup.
-    """
-    pincode_data = load_maharashtra_pincode_data()
-    return {entry["pincode"]: entry for entry in pincode_data if "pincode" in entry}
-
-
-@lru_cache(maxsize=1)
-def _load_maharashtra_mla_map() -> Dict[str, Dict[str, Any]]:
-    """
-    Load and cache Maharashtra MLA information as a dict for O(1) lookup.
-    """
-    mla_data = load_maharashtra_mla_data()
-    return {entry["assembly_constituency"]: entry for entry in mla_data if "assembly_constituency" in entry}
-
-
 def find_constituency_by_pincode(pincode: str) -> Optional[Dict[str, Any]]:
     """
     Find constituency information by pincode.
@@ -135,6 +106,7 @@ def find_constituency_by_pincode(pincode: str) -> Optional[Dict[str, Any]]:
     pincode_map = _load_maharashtra_pincode_map()
     entry = pincode_map.get(pincode)
 
+    entry = pincode_data.get(pincode)
     if entry:
         return {
             "district": entry.get("district"),
@@ -162,6 +134,7 @@ def find_mla_by_constituency(constituency_name: str) -> Optional[Dict[str, Any]]
     mla_map = _load_maharashtra_mla_map()
     entry = mla_map.get(constituency_name)
 
+    entry = mla_data.get(constituency_name)
     if entry:
         return {
             "mla_name": entry.get("mla_name"),

tests/test_mh_endpoint.py

@@ -35,6 +35,7 @@ def test_maharashtra_endpoint():
         print(f"   Party: {data['mla'].get('party')}")
         assert data['pincode'] == '411001'
         assert data['district'] == 'Pune'
+        # The data was updated to real world data
         assert data['mla']['name'] == 'Ravindra Dhangekar'
     else:
         print(f"   ✗ Failed: {response.json()}")