Update web_search.py

web_search.py

@@ -1,17 +1,12 @@
-# web_search.py
+# ====================================================
+# web_search.py — Village Enabled
+# ====================================================
+
 """
 Minimal Tavily wrapper: run a web search and return a list of UDISE codes.
 
-Usage:
-    from web_search import tavily_search_codes
-    codes = tavily_search_codes("GOVT HIGH SCHOOL", state_name="Karnataka", district="Bengaluru")
-
-Behavior:
-    - Uses TAVILY_API_KEY env var if api_key is not provided.
-    - Extracts 11-digit sequences that look like UDISE codes.
-    - If state_name provided and matches a known mapping, only returns codes whose
-      first two digits match that state's prefix.
-    - Returns a list of unique UDISE codes (strings) in the order they were found.
+Enhancement:
+- Optional village support for more precise queries
 """
 
 import os
@@ -24,46 +19,62 @@ try:
 except Exception:
     TavilyClient = None
 
-# Minimal mapping of state name -> 2-digit UDISE prefix.
-# Extend this map to match your canonical state names if you want strict filtering.
+# ----------------------------------------------------
+# State → UDISE prefix mapping
+# ----------------------------------------------------
 STATE_TO_UDISE_CODE = {
     "Jammu & Kashmir": "01", "Himachal Pradesh": "02", "Punjab": "03",
     "Chandigarh": "04", "Uttarakhand": "05", "Haryana": "06", "Delhi": "07",
     "Rajasthan": "08", "Uttar Pradesh": "09", "Bihar": "10", "Sikkim": "11",
     "Arunachal Pradesh": "12", "Nagaland": "13", "Manipur": "14", "Mizoram": "15",
     "Tripura": "16", "Meghalaya": "17", "Assam": "18", "West Bengal": "19",
-    "Jharkhand": "20", "Odisha": "21", "Chhattisgarh": "22", "Madhya Pradesh": "23",
-    "Gujarat": "24", "Daman & Diu": "25", "Dadra & Nagar Haveli": "26",
-    "Maharashtra": "27", "Andhra Pradesh": "28", "Karnataka": "29",
-    "Goa": "30", "Lakshadweep": "31", "Kerala": "32", "Tamil Nadu": "33",
-    "Puducherry": "34", "Andaman & Nicobar Islands": "35", "Telangana": "36",
-    "Ladakh": "37"
+    "Jharkhand": "20", "Odisha": "21", "Chhattisgarh": "22",
+    "Madhya Pradesh": "23", "Gujarat": "24", "Daman & Diu": "25",
+    "Dadra & Nagar Haveli": "26", "Maharashtra": "27",
+    "Andhra Pradesh": "28", "Karnataka": "29", "Goa": "30",
+    "Lakshadweep": "31", "Kerala": "32", "Tamil Nadu": "33",
+    "Puducherry": "34", "Andaman & Nicobar Islands": "35",
+    "Telangana": "36", "Ladakh": "37",
 }
 
-# compile once: strict 11-digit match
+# strict 11-digit UDISE match
 _UDISE_RE = re.compile(r"(?<!\d)(\d{11})(?!\d)")
 
-def _build_query(school_name: Optional[str], state_name: Optional[str], district: Optional[str]) -> str:
-    parts = []
-    parts.append("UDISE code of")
+
+# ----------------------------------------------------
+# Query builder (Village-aware)
+# ----------------------------------------------------
+def _build_query(
+    school_name: Optional[str],
+    state_name: Optional[str],
+    district: Optional[str],
+    village: Optional[str] = None,
+) -> str:
+    parts = ["UDISE code of"]
+
     if school_name:
         parts.append(f"School {school_name.strip()}")
-    if state_name:
-        parts.append(f"in state {state_name.strip()}")
+
+    if village:
+        parts.append(f"in village {village.strip()}")
+
     if district:
-        parts.append(f"of district {district.strip()}")
+        parts.append(f"district {district.strip()}")
+
+    if state_name:
+        parts.append(f"state {state_name.strip()}")
+
     return " ".join(parts).strip()
 
 
 def _call_tavily(api_key: Optional[str], query: str):
     key = api_key or os.getenv("TAVILY_API_KEY")
     if not key:
-        return {"ok": False, "error": "No Tavily API key provided (set TAVILY_API_KEY or pass api_key)."}
+        return {"ok": False, "error": "No Tavily API key provided."}
     if TavilyClient is None:
-        return {"ok": False, "error": "tavily package not installed or TavilyClient not importable."}
+        return {"ok": False, "error": "tavily package not installed."}
     try:
         client = TavilyClient(key)
-        # SDK may return complex object; keep raw
         resp = client.search(query=query, country="india")
         return {"ok": True, "data": resp}
     except Exception as e:
@@ -73,11 +84,7 @@ def _call_tavily(api_key: Optional[str], query: str):
 def _normalize_state_key(state_name: Optional[str]) -> Optional[str]:
     if not state_name:
         return None
-    # simple exact / title / upper checks
-    for k in (state_name, state_name.title(), state_name.upper()):
-        if k in STATE_TO_UDISE_CODE:
-            return k
-    # fallback: compare cleaned names (remove non-letters)
+
     cleaned = re.sub(r"[^A-Za-z]", "", state_name).lower()
     for k in STATE_TO_UDISE_CODE:
         if re.sub(r"[^A-Za-z]", "", k).lower() == cleaned:
@@ -85,81 +92,80 @@ def _normalize_state_key(state_name: Optional[str]) -> Optional[str]:
     return None
 
 
+# ----------------------------------------------------
+# Public API
+# ----------------------------------------------------
 def tavily_search_codes(
     school_name: Optional[str],
     state_name: Optional[str] = None,
     district: Optional[str] = None,
+    village: Optional[str] = None,
     api_key: Optional[str] = None,
-    enforce_state_prefix: bool = True
+    enforce_state_prefix: bool = True,
 ) -> List[str]:
     """
-    Perform a Tavily search and return a list of unique UDISE codes (strings).
-
-    Args:
-        school_name: partial or full school name to search
-        state_name: optional state name to restrict to that state's UDISE prefix
-        district: optional district to include in query
-        api_key: optional Tavily API key (falls back to TAVILY_API_KEY env)
-        enforce_state_prefix: if True and state_name is known, only return codes
-                              whose first two digits match that state's prefix.
-
-    Returns:
-        List[str] of unique UDISE codes in the order found.
+    Perform a Tavily search and return a list of unique UDISE codes.
+
+    Village is used only to improve search precision.
     """
+
     if not school_name:
-        return None
-    query = _build_query(school_name, state_name, district)
+        return []
+
+    query = _build_query(school_name, state_name, district, village)
     call = _call_tavily(api_key, query)
+
     if not call.get("ok"):
-        # on error return empty list (caller can inspect logs)
         return []
 
     raw = call.get("data") or {}
-    # flatten likely text fields from raw into strings to search for codes
     snippets = []
 
-    # common shapes: dict with 'results' or 'data' or a list
     if isinstance(raw, dict):
-        # try results/data/results/data.items etc
-        candidates = raw.get("results") or raw.get("data", {}).get("results") or raw.get("data") or raw.get("hits") or raw.get("items")
-        if isinstance(candidates, list):
-            for item in candidates:
-                if isinstance(item, dict):
-                    snippets.append(" ".join([str(item.get("title", "")), str(item.get("content", "")), str(item.get("text", "")), str(item.get("url", ""))]))
-                else:
-                    snippets.append(str(item))
-        else:
-            # fallback stringify the whole dict
-            snippets.append(str(raw))
+        candidates = (
+            raw.get("results")
+            or raw.get("data", {}).get("results")
+            or raw.get("items")
+            or []
+        )
+        for item in candidates:
+            if isinstance(item, dict):
+                snippets.append(
+                    " ".join(
+                        [
+                            str(item.get("title", "")),
+                            str(item.get("content", "")),
+                            str(item.get("text", "")),
+                            str(item.get("url", "")),
+                        ]
+                    )
+                )
+            else:
+                snippets.append(str(item))
     elif isinstance(raw, list):
-        for item in raw:
-            snippets.append(str(item))
+        snippets = [str(x) for x in raw]
     else:
-        snippets.append(str(raw))
+        snippets = [str(raw)]
 
-    # Determine allowed prefix if requested
     allowed_prefix = None
     state_key = _normalize_state_key(state_name)
     if enforce_state_prefix and state_key:
         allowed_prefix = STATE_TO_UDISE_CODE.get(state_key)
 
-    found = []
-    seen = set()
+    found, seen = [], set()
     for text in snippets:
-        if not text:
-            continue
         for m in _UDISE_RE.finditer(text):
             code = m.group(1)
-            if not code or code in seen:
+            if code in seen:
                 continue
-            # quick prefix check
+
             prefix = code[:2]
-            # ignore codes that do not start with a valid state prefix
-            if prefix not in set(STATE_TO_UDISE_CODE.values()):
+            if prefix not in STATE_TO_UDISE_CODE.values():
                 continue
             if allowed_prefix and prefix != allowed_prefix:
                 continue
+
             seen.add(code)
             found.append(code)
 
-    return found
+    return found