Fix destination extraction and remove BKK default

app/service.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-import json
 import logging
 import os
 import random
 import re
 import time
 import uuid
-from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Tuple
 
 from pydantic import BaseModel, Field
@@ -15,17 +13,10 @@ from pydantic import BaseModel, Field
 LOGGER = logging.getLogger("flight-agent-api")
 logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO").upper())
 
-DEFAULT_BUCKET_PATH = "/mnt/app-storage"
-DEFAULT_LOCAL_FALLBACK = "/tmp/flight-agent-data"
-
-
-# ============================================================
-# Models
-# ============================================================
 
 class TripRequest(BaseModel):
     origin: Optional[str] = None
-    destination: str = "BKK"
+    destination: Optional[str] = None
     depart_date: Optional[str] = None
     return_date: Optional[str] = None
     adults: int = 1
@@ -90,91 +81,6 @@ class AgentCoreState(BaseModel):
         self.audit.append({"ts": time.time(), "kind": kind, "payload": payload})
 
 
-# ============================================================
-# Storage
-# ============================================================
-
-
-def _sanitize_thread_id(thread_id: str) -> str:
-    value = re.sub(r"[^A-Za-z0-9_.-]", "_", thread_id or "anonymous")
-    return value[:120] or "anonymous"
-
-
-class StorageManager:
-    def __init__(self) -> None:
-        requested = os.getenv("APP_STORAGE_PATH", DEFAULT_BUCKET_PATH).strip() or DEFAULT_BUCKET_PATH
-        fallback = os.getenv("LOCAL_STORAGE_FALLBACK", DEFAULT_LOCAL_FALLBACK).strip() or DEFAULT_LOCAL_FALLBACK
-        path, mode = self._resolve_storage_path(requested, fallback)
-        self.base_path = path
-        self.mode = mode
-        self.sessions_dir = self.base_path / "sessions"
-        self.logs_dir = self.base_path / "logs"
-        self.cache_dir = self.base_path / "cache"
-        for folder in (self.sessions_dir, self.logs_dir, self.cache_dir):
-            folder.mkdir(parents=True, exist_ok=True)
-        LOGGER.info("Storage initialized at %s (mode=%s)", self.base_path, self.mode)
-
-    def _resolve_storage_path(self, requested: str, fallback: str) -> Tuple[Path, str]:
-        requested_path = Path(requested)
-        try:
-            requested_path.mkdir(parents=True, exist_ok=True)
-            probe = requested_path / ".write_test"
-            probe.write_text("ok", encoding="utf-8")
-            probe.unlink(missing_ok=True)
-            if requested_path == Path(DEFAULT_BUCKET_PATH):
-                return requested_path, "bucket"
-            return requested_path, "custom"
-        except Exception as exc:  # pragma: no cover - defensive
-            LOGGER.warning("Bucket/custom storage path unavailable (%s). Falling back to %s", exc, fallback)
-            fallback_path = Path(fallback)
-            fallback_path.mkdir(parents=True, exist_ok=True)
-            return fallback_path, "local-fallback"
-
-    def get_info(self) -> Dict[str, Any]:
-        return {
-            "base_path": str(self.base_path),
-            "mode": self.mode,
-            "sessions_dir": str(self.sessions_dir),
-            "logs_dir": str(self.logs_dir),
-            "cache_dir": str(self.cache_dir),
-        }
-
-    def _session_file(self, thread_id: str) -> Path:
-        return self.sessions_dir / f"{_sanitize_thread_id(thread_id)}.json"
-
-    def load_session(self, thread_id: str) -> Optional[AgentCoreState]:
-        path = self._session_file(thread_id)
-        if not path.exists():
-            return None
-        try:
-            data = json.loads(path.read_text(encoding="utf-8"))
-            return AgentCoreState.model_validate(data)
-        except Exception as exc:  # pragma: no cover - defensive
-            LOGGER.warning("Failed to load session %s: %s", thread_id, exc)
-            return None
-
-    def save_session(self, thread_id: str, state: AgentCoreState) -> None:
-        path = self._session_file(thread_id)
-        path.write_text(json.dumps(state.model_dump(), ensure_ascii=False, indent=2), encoding="utf-8")
-
-    def delete_session(self, thread_id: str) -> None:
-        self._session_file(thread_id).unlink(missing_ok=True)
-
-    def append_event(self, thread_id: str, kind: str, payload: Dict[str, Any]) -> None:
-        event = {
-            "ts": time.time(),
-            "thread_id": thread_id,
-            "kind": kind,
-            "payload": payload,
-        }
-        with (self.logs_dir / "events.jsonl").open("a", encoding="utf-8") as f:
-            f.write(json.dumps(event, ensure_ascii=False) + "\n")
-
-
-# ============================================================
-# Guardrails + NLU helpers
-# ============================================================
-
 PII_PATTERNS: List[Tuple[re.Pattern, str]] = [
     (re.compile(r"\b\d{12,19}\b"), "[REDACTED_CARD]"),
     (re.compile(r"\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b", re.I), "[REDACTED_EMAIL]"),
@@ -196,7 +102,29 @@ STOPS_RE = re.compile(r"\b(nonstop|direct|0\s*stop|1\s*stop|2\s*stop)\b", re.I)
 ADULTS_RE = re.compile(r"\b(\d+)\s*(adult|adults|passenger|passengers|traveler|travelers)\b", re.I)
 
 SUPPORT_WORDS = {"refund", "cancellation", "cancel", "change fee", "baggage", "reschedule", "policy", "rules", "luggage"}
-BOOKING_WORDS = {"flight", "book", "ticket", "bangkok", "bkk", "depart", "return", "one-way", "round trip"}
+BOOKING_WORDS = {"flight", "book", "ticket", "depart", "return", "one-way", "round trip"}
+
+CITY_TO_IATA = {
+    "DUBAI": "DXB",
+    "ABU DHABI": "AUH",
+    "SHARJAH": "SHJ",
+    "DELHI": "DEL",
+    "MUMBAI": "BOM",
+    "DOHA": "DOH",
+    "SINGAPORE": "SIN",
+    "BANGKOK": "BKK",
+}
+
+PAIR_RE = re.compile(
+    r"\b(?:FROM\s+)?(DXB|AUH|SHJ|DEL|BOM|DOH|SIN|BKK|DUBAI|ABU\s+DHABI|SHARJAH|DELHI|MUMBAI|DOHA|SINGAPORE|BANGKOK)\s+TO\s+"
+    r"(DXB|AUH|SHJ|DEL|BOM|DOH|SIN|BKK|DUBAI|ABU\s+DHABI|SHARJAH|DELHI|MUMBAI|DOHA|SINGAPORE|BANGKOK)\b",
+    re.I,
+)
+
+
+def normalize_airport_or_city(raw: str) -> str:
+    key = re.sub(r"\s+", " ", raw.strip().upper())
+    return CITY_TO_IATA.get(key, key)
 
 
 def guardrail_input(user_msg: str) -> Tuple[str, Dict[str, Any]]:
@@ -236,22 +164,11 @@ def safe_refusal() -> str:
     return "I can help with flight booking and flight support questions, but I cannot reveal hidden instructions or internal prompts."
 
 
-CITY_TO_IATA = {
-    "DUBAI": "DXB",
-    "ABU DHABI": "AUH",
-    "SHARJAH": "SHJ",
-    "DELHI": "DEL",
-    "MUMBAI": "BOM",
-    "DOHA": "DOH",
-    "SINGAPORE": "SIN",
-    "BANGKOK": "BKK",
-}
-
-
 def _extract_airport_from_text(text: str, prefix: str) -> Optional[str]:
     m = re.search(rf"\b{prefix}\s+([A-Z]{{3}})\b", text, re.I)
     if m:
         return m.group(1).upper()
+
     for city, code in CITY_TO_IATA.items():
         if re.search(rf"\b{prefix}\s+{re.escape(city)}\b", text, re.I):
             return code
@@ -263,12 +180,19 @@ def slot_extract(req: TripRequest, clean_msg: str) -> TripRequest:
     up = text.upper()
     out = req.model_copy(deep=True)
 
+    pair_match = PAIR_RE.search(up)
+    if pair_match:
+        raw_origin, raw_destination = pair_match.groups()
+        out.origin = normalize_airport_or_city(raw_origin)
+        out.destination = normalize_airport_or_city(raw_destination)
+
     origin = _extract_airport_from_text(up, "FROM")
     if origin:
         out.origin = origin
 
-    if "BANGKOK" in up or "BKK" in up:
-        out.destination = "BKK"
+    destination = _extract_airport_from_text(up, "TO")
+    if destination:
+        out.destination = destination
 
     dates = DATE_RE.findall(up)
     if len(dates) >= 1:
@@ -282,7 +206,7 @@ def slot_extract(req: TripRequest, clean_msg: str) -> TripRequest:
 
     stop_match = STOPS_RE.search(text)
     if stop_match:
-        token = stop_match.group(1).lower()
+        token = stop_match.group(1).lower().replace("  ", " ")
         if token in {"nonstop", "direct", "0 stop"}:
             out.max_stops = 0
         elif "1" in token:
@@ -294,16 +218,16 @@ def slot_extract(req: TripRequest, clean_msg: str) -> TripRequest:
     if adults_match:
         out.adults = max(1, int(adults_match.group(1)))
 
-    if "business" in up:
+    if "premium economy" in up:
+        out.cabin = "PREMIUM_ECONOMY"
+    elif "business" in up:
         out.cabin = "BUSINESS"
     elif "first" in up:
         out.cabin = "FIRST"
-    elif "premium economy" in up:
-        out.cabin = "PREMIUM_ECONOMY"
     elif "economy" in up:
         out.cabin = "ECONOMY"
 
-    if "checked" in up or "check-in" in up:
+    if "checked" in up or "CHECK-IN" in up:
         out.baggage = "CHECKED"
     elif "no baggage" in up:
         out.baggage = "NONE"
@@ -322,7 +246,13 @@ def heuristic_intent(clean_msg: str) -> Tuple[str, str, float]:
     if any(k in m for k in SUPPORT_WORDS):
         return "TRAVEL_SUPPORT", "POLICY_EXPLAIN", 0.88
 
-    if any(k in m for k in BOOKING_WORDS) or DATE_RE.search(m) or _extract_airport_from_text(m.upper(), "FROM"):
+    if (
+        any(k in m for k in BOOKING_WORDS)
+        or DATE_RE.search(m)
+        or _extract_airport_from_text(m.upper(), "FROM")
+        or _extract_airport_from_text(m.upper(), "TO")
+        or PAIR_RE.search(m.upper())
+    ):
         return "TRAVEL_BOOKING", "FLIGHT_SEARCH", 0.84
 
     if any(k in m for k in ["hi", "hello", "hey"]):
@@ -335,23 +265,24 @@ def missing_slots(req: TripRequest) -> List[str]:
     miss: List[str] = []
     if not req.origin:
         miss.append("origin (for example DXB)")
+    if not req.destination:
+        miss.append("destination (for example BKK or Bangkok)")
     if not req.depart_date:
         miss.append("depart date in YYYY-MM-DD")
     return miss
 
 
-# ============================================================
-# Mock tools
-# ============================================================
-
 CARRIERS = ["EK", "QR", "EY", "TG", "SQ", "AI"]
+HUBS = ["DOH", "SIN", "KUL", "BOM"]
 
 
 def make_mock_flight_options(req: TripRequest, k: int = 10) -> List[FlightOption]:
-    if not req.origin or not req.depart_date:
+    if not req.origin or not req.destination or not req.depart_date:
         return []
 
-    rng = random.Random(f"{req.origin}-{req.depart_date}-{req.max_price_usd}-{req.max_stops}-{req.cabin}")
+    rng = random.Random(
+        f"{req.origin}-{req.destination}-{req.depart_date}-{req.max_price_usd}-{req.max_stops}-{req.cabin}"
+    )
     options: List[FlightOption] = []
 
     for _ in range(k):
@@ -367,7 +298,7 @@ def make_mock_flight_options(req: TripRequest, k: int = 10) -> List[FlightOption
             legs.append(
                 FlightLeg(
                     from_airport=req.origin,
-                    to_airport="BKK",
+                    to_airport=req.destination,
                     depart_time=f"{req.depart_date}T09:{rng.randint(10, 59):02d}",
                     arrive_time=f"{req.depart_date}T{rng.randint(14, 20):02d}:{rng.randint(10, 59):02d}",
                     carrier=carrier,
@@ -375,7 +306,7 @@ def make_mock_flight_options(req: TripRequest, k: int = 10) -> List[FlightOption
                 )
             )
         else:
-            hub = rng.choice(["DOH", "SIN", "KUL", "BOM"])
+            hub = rng.choice([h for h in HUBS if h not in {req.origin, req.destination}] or HUBS)
             legs.append(
                 FlightLeg(
                     from_airport=req.origin,
@@ -389,7 +320,7 @@ def make_mock_flight_options(req: TripRequest, k: int = 10) -> List[FlightOption
             legs.append(
                 FlightLeg(
                     from_airport=hub,
-                    to_airport="BKK",
+                    to_airport=req.destination,
                     depart_time=f"{req.depart_date}T{rng.randint(12, 16):02d}:{rng.randint(10, 59):02d}",
                     arrive_time=f"{req.depart_date}T{rng.randint(16, 23):02d}:{rng.randint(10, 59):02d}",
                     carrier=carrier,
@@ -452,14 +383,19 @@ def send_email(email: str, subject: str, body: str) -> Dict[str, Any]:
     }
 
 
-# ============================================================
-# Ranking + formatting
-# ============================================================
-
 POLICY_DOCS = [
-    {"source_id": "policy_refund_1", "text": "Refund eligibility depends on fare brand. FLEX fares generally have lower change and cancellation penalties than LITE fares."},
-    {"source_id": "policy_baggage_1", "text": "Checked baggage inclusion depends on fare brand and route. Cabin baggage is usually included unless the fare is very restrictive."},
-    {"source_id": "policy_change_1", "text": "Changes may be permitted before departure subject to fare rules, price difference, and applicable service fees."},
+    {
+        "source_id": "policy_refund_1",
+        "text": "Refund eligibility depends on fare brand. FLEX fares generally have lower change and cancellation penalties than LITE fares.",
+    },
+    {
+        "source_id": "policy_baggage_1",
+        "text": "Checked baggage inclusion depends on fare brand and route. Cabin baggage is usually included unless the fare is very restrictive.",
+    },
+    {
+        "source_id": "policy_change_1",
+        "text": "Changes may be permitted before departure subject to fare rules, price difference, and applicable service fees.",
+    },
 ]
 
 
@@ -506,37 +442,17 @@ def present_topk(state: AgentCoreState, options: List[FlightOption], k: int = 3)
     return "\n".join(lines)
 
 
-# ============================================================
-# Agent service
-# ============================================================
-
 class FlightBookingAgentService:
     def __init__(self) -> None:
         self.sessions: Dict[str, AgentCoreState] = {}
-        self.storage = StorageManager()
-
-    def get_storage_info(self) -> Dict[str, Any]:
-        return self.storage.get_info()
 
     def _get_state(self, thread_id: str) -> AgentCoreState:
-        if thread_id in self.sessions:
-            return self.sessions[thread_id]
-        loaded = self.storage.load_session(thread_id)
-        if loaded is not None:
-            self.sessions[thread_id] = loaded
-            return loaded
-        self.sessions[thread_id] = AgentCoreState()
+        if thread_id not in self.sessions:
+            self.sessions[thread_id] = AgentCoreState()
         return self.sessions[thread_id]
 
-    def _persist(self, thread_id: str, state: AgentCoreState, event_kind: Optional[str] = None) -> None:
-        self.storage.save_session(thread_id, state)
-        if event_kind:
-            self.storage.append_event(thread_id, event_kind, {"audit_count": len(state.audit)})
-
     def reset_state(self, thread_id: str) -> None:
         self.sessions.pop(thread_id, None)
-        self.storage.delete_session(thread_id)
-        self.storage.append_event(thread_id, "session.reset", {})
 
     def get_state(self, thread_id: str) -> Dict[str, Any]:
         return self._get_state(thread_id).model_dump()
@@ -554,7 +470,7 @@ class FlightBookingAgentService:
         if meta.get("injection_suspected"):
             response = safe_refusal()
             core.log("refuse", {"reason": "prompt_injection"})
-            return self._response(thread_id, response, core, "chat.refusal")
+            return self._response(thread_id, response, core)
 
         parent, child, conf = heuristic_intent(clean_msg)
         if parent == "TRAVEL_BOOKING":
@@ -565,44 +481,44 @@ class FlightBookingAgentService:
 
         if core.booking_hold and core.booking_hold.status == "HELD" and core.user_contact.get("email"):
             response = self._confirm_booking(core)
-            return self._response(thread_id, response, core, "chat.auto_confirm")
+            return self._response(thread_id, response, core)
 
         if parent == "TRAVEL_SUPPORT":
             response = self._support(clean_msg, core)
-            return self._response(thread_id, response, core, "chat.support")
+            return self._response(thread_id, response, core)
 
         if parent == "TRAVEL_BOOKING" and child == "FLIGHT_BOOK_HOLD":
             response = self._create_hold(clean_msg, core)
-            return self._response(thread_id, response, core, "chat.hold")
+            return self._response(thread_id, response, core)
 
         if parent == "TRAVEL_BOOKING":
             response = self._search_flights(core)
-            return self._response(thread_id, response, core, "chat.search")
+            return self._response(thread_id, response, core)
 
         if parent == "SMALLTALK":
-            response = "I can help book flights to Bangkok. Try: From DXB on 2026-07-10 under 650"
-            return self._response(thread_id, response, core, "chat.smalltalk")
+            response = "I can help book flights. Try: From DXB to BKK on 2026-07-10 under 650"
+            return self._response(thread_id, response, core)
 
         response = "I can only help with flight booking or flight support in this demo."
-        return self._response(thread_id, response, core, "chat.out_of_scope")
+        return self._response(thread_id, response, core)
 
     def search(self, thread_id: str, req: TripRequest) -> Dict[str, Any]:
         core = self._get_state(thread_id)
         core.request = req
         response = self._search_flights(core)
-        return self._response(thread_id, response, core, "api.search")
+        return self._response(thread_id, response, core)
 
     def hold(self, thread_id: str, choice_index: int) -> Dict[str, Any]:
         core = self._get_state(thread_id)
         clean_msg = f"book {choice_index}"
         response = self._create_hold(clean_msg, core)
-        return self._response(thread_id, response, core, "api.hold")
+        return self._response(thread_id, response, core)
 
     def confirm(self, thread_id: str, email: str) -> Dict[str, Any]:
         core = self._get_state(thread_id)
         core.user_contact["email"] = email
         response = self._confirm_booking(core)
-        return self._response(thread_id, response, core, "api.confirm")
+        return self._response(thread_id, response, core)
 
     def _support(self, clean_msg: str, core: AgentCoreState) -> str:
         docs = retrieve_policy_notes(clean_msg, topk=2)
@@ -615,8 +531,8 @@ class FlightBookingAgentService:
         miss = missing_slots(core.request)
         if miss:
             msg = (
-                f"I can help book your flight to Bangkok (BKK). Please provide: {', '.join(miss)}.\n"
-                f"Example: From DXB on 2026-07-10 under 650"
+                f"I can help book your flight. Please provide: {', '.join(miss)}.\n"
+                f"Example: From DXB to BKK on 2026-07-10 under 650"
             )
             core.log("search.ask_clarify", {"missing": miss})
             return msg
@@ -697,12 +613,9 @@ class FlightBookingAgentService:
         core.log("tool.send_email", email_res)
         return f"Booking confirmed for {email}. Confirmation reference: {core.booking_hold.hold_id}."
 
-    def _response(self, thread_id: str, response: str, core: AgentCoreState, event_kind: str) -> Dict[str, Any]:
-        payload = {
+    def _response(self, thread_id: str, response: str, core: AgentCoreState) -> Dict[str, Any]:
+        return {
             "thread_id": thread_id,
             "response": response,
             "state": core.model_dump(),
-            "storage": self.get_storage_info(),
         }
-        self._persist(thread_id, core, event_kind)
-        return payload