forked from zhouruizhe/gmTouringMiniApp
106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
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
|