planmate/llm.py

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