Create web_search.py

web_search.py

@@ -0,0 +1,165 @@
+# web_search.py
+"""
+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.
+"""
+
+import os
+import re
+from typing import List, Optional
+
+# optional Tavily SDK
+try:
+    from tavily import TavilyClient
+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_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"
+}
+
+# compile once: strict 11-digit 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 = []
+    if school_name:
+        parts.append(school_name.strip())
+    if district:
+        parts.append(f"district {district.strip()}")
+    if state_name:
+        parts.append(f"state {state_name.strip()}")
+    # bias to India and ask specifically for UDISE mentions
+    parts.append("India")
+    parts.append("UDISE code OR UDISE number OR Udise")
+    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)."}
+    if TavilyClient is None:
+        return {"ok": False, "error": "tavily package not installed or TavilyClient not importable."}
+    try:
+        client = TavilyClient(key)
+        # SDK may return complex object; keep raw
+        resp = client.search(query=query)
+        return {"ok": True, "data": resp}
+    except Exception as e:
+        return {"ok": False, "error": str(e)}
+
+
+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:
+            return k
+    return None
+
+
+def tavily_search_codes(
+    school_name: Optional[str],
+    state_name: Optional[str] = None,
+    district: Optional[str] = None,
+    api_key: Optional[str] = None,
+    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.
+    """
+    query = _build_query(school_name, state_name, district)
+    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))
+    elif isinstance(raw, list):
+        for item in raw:
+            snippets.append(str(item))
+    else:
+        snippets.append(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()
+    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:
+                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()):
+                continue
+            if allowed_prefix and prefix != allowed_prefix:
+                continue
+            seen.add(code)
+            found.append(code)
+
+    return found