from __future__ import annotations import json from pathlib import Path from typing import Any, Optional from .schemas import TravelPreferences DATA_DIR = Path(__file__).resolve().parent.parent / "data" class KnowledgeBase: def __init__(self) -> None: self.places: list[dict[str, Any]] = json.loads( (DATA_DIR / "places.json").read_text(encoding="utf-8") ) travel_rows = json.loads( (DATA_DIR / "travel_times.json").read_text(encoding="utf-8") ) self.travel_times = { (row["from"], row["to"]): row for row in travel_rows } @property def count(self) -> int: return len(self.places) def by_id(self, place_id: str) -> Optional[dict[str, Any]]: return next((place for place in self.places if place["id"] == place_id), None) def retrieve(self, preferences: TravelPreferences, limit: int = 8) -> list[dict[str, Any]]: requirements = preferences.extra_requirements.lower() indoor_required = any( token in requirements for token in ("室内", "下雨", "雨天", "避雨") ) avoid_climbing = any( token in requirements for token in ("不爬山", "不要爬山", "不登山", "体力不好", "婴儿车") ) youngest = min(preferences.child_ages) if preferences.child_ages else None scored: list[tuple[int, str, dict[str, Any]]] = [] for place in self.places: if place["city"] != "深圳市" or place["district"] != "光明区": continue if youngest is not None and youngest < place.get("minAge", 0): continue if avoid_climbing and place.get("intensity") == "high": continue if indoor_required and not ( place.get("indoor") or place.get("rainyDaySuitable") ): continue score = 0 score += 3 * len(set(preferences.themes) & set(place.get("themes", []))) score += 3 * len( set(preferences.interests) & set(place.get("interests", [])) ) if preferences.children and "儿童" in place.get("suitableFor", []): score += 2 elif not preferences.children and "成人" in place.get("suitableFor", []): score += 2 if preferences.pace.value in place.get("pace", []): score += 1 if preferences.budget_level.value == place.get("budgetLevel"): score += 1 if indoor_required and ( place.get("indoor") or place.get("rainyDaySuitable") ): score += 2 scored.append((score, place["name"], place)) scored.sort(key=lambda item: (-item[0], item[1])) return [place for _, _, place in scored[:limit]] def travel_minutes( self, from_id: str, to_id: str, transport: str ) -> Optional[int]: row = self.travel_times.get((from_id, to_id)) or self.travel_times.get( (to_id, from_id) ) if not row: return None value = row.get(transport) return int(value) if value is not None else None def travel_context( self, candidates: list[dict[str, Any]], transport: str ) -> list[dict[str, Any]]: ids = {place["id"] for place in candidates} result = [] for row in self.travel_times.values(): if row["from"] in ids and row["to"] in ids: result.append( { "from": row["from"], "to": row["to"], "minutes": row.get(transport), "note": row.get("note", "POC估算,请以实际导航为准"), } ) return result