import json from typing import List, Dict, Optional, Tuple import google.generativeai as genai from .config import get_gemini_key from .utils import is_iata_code # Initialize once _gen_inited = False def _ensure_init(): global _gen_inited if not _gen_inited: genai.configure(api_key=get_gemini_key()) _gen_inited = True def _model(): _ensure_init() return genai.GenerativeModel( model_name="gemini-2.0-flash-lite", generation_config={ "temperature": 0.2, "top_p": 0.9, "top_k": 40, "max_output_tokens": 600, }, ) def resolve_city_to_iata_ai(city: str, country_hint: Optional[str] = None) -> Tuple[str, str, str]: """Use the LLM to resolve a human-readable city to an IATA code. Returns (code, canonical_name, kind) where kind is 'CITY' or 'AIRPORT'. Raises RuntimeError on failure. """ m = _model() prompt = f""" You convert a city name into a 3-letter IATA code for flight search. Rules: - If the metro area has multiple major airports with an aggregate IATA city code, return the CITY code (e.g., New York→NYC, London→LON, Paris→PAR, Tokyo→TYO). - If the city typically uses a single primary commercial airport, return that AIRPORT code (e.g., Dubai→DXB, Doha→DOH, Lahore→LHE, Karachi→KHI, Istanbul→IST). - Return empty code "" if no suitable airport exists. - The code MUST be exactly 3 uppercase letters A–Z. - Respond with ONLY a compact JSON object, no additional text, using this schema: {{"code":"XXX","name":"Canonical City or Airport Name","kind":"CITY|AIRPORT","alternates":["AAA","BBB"]}} Input: - city: {city} - country_hint: {country_hint or 'unknown'} """ resp = m.generate_content(prompt) text = (resp.text or "").strip() # Extract JSON block mjson = None if "{" in text: try: start = text.find("{") end = text.rfind("}") + 1 mjson = json.loads(text[start:end]) except Exception: pass if not mjson or not isinstance(mjson, dict): raise RuntimeError("AI could not produce a valid JSON mapping for the city.") code = str(mjson.get("code", "")).strip().upper() name = str(mjson.get("name", city)).strip() kind = str(mjson.get("kind", "")).strip().upper() if not is_iata_code(code): raise RuntimeError(f"AI returned an invalid IATA code: {code}") if kind not in ("CITY", "AIRPORT"): kind = "AIRPORT" return code, name, kind def rank_accommodations(accommodations: List[Dict], prefs: str = "") -> List[Dict]: """Add an 'llm_score' to each accommodation and sort by it.""" if not accommodations: return [] m = _model() lines = "\n".join( [ f"- {a.get('name','(no name)')} | rate:{a.get('rate')} | dist:{int(a.get('dist',0))}m | kinds:{a.get('kinds','')}" for a in accommodations[:50] ] ) prompt = f"""Rank the following accommodations for a tourist trip. Prefer central, well-rated options. User preferences: {prefs or 'not specified'}. Return a JSON array of objects with 'name' and an integer 'score' 1-100. Items: {lines} """ resp = m.generate_content(prompt) try: text = resp.text or "" if "```json" in text: data = json.loads(text.split("```json")[-1].split("```")[0]) else: data = json.loads(text) except Exception: data = [] score_map = {d.get("name", ""): int(d.get("score", 50)) for d in data if isinstance(d, dict)} for a in accommodations: a["llm_score"] = score_map.get(a.get("name", ""), 50) return sorted(accommodations, key=lambda x: x.get("llm_score", 50), reverse=True) def generate_itinerary( city: str, start_date: str, days: int, selected_attractions: List[Dict], selected_stay: Dict | None, weather_summary: str, ) -> str: m = _model() attractions_text = "\n".join([f"- {a.get('name')} ({a.get('kinds','')})" for a in selected_attractions]) hotel_text = selected_stay.get("name") if selected_stay else "TBD" prompt = f""" Create a practical, day-by-day itinerary for a trip. City: {city} Start Date: {start_date} Days: {days} Hotel: {hotel_text} Weather (summary): {weather_summary} Attractions to consider: {attractions_text} Constraints: group nearby sights, account for weather, add meal/time suggestions, include commute notes. Return Markdown with sections Day 1, Day 2, ..., and a brief daily plan (morning/afternoon/evening). """ resp = m.generate_content(prompt) return resp.text