"""Retrieval grounding module for fetching similar game examples. Handles four dataset schemas: 1. Old dataset (games_dataset.json) — uses ``expected_output`` key 2. New scavenger_hunt dataset — uses ``output`` key with ``tasks`` array 3. New hide_and_seek dataset — uses ``output`` key with ``hiding_zones`` + ``play_area`` 4. New tag dataset — uses ``output`` key with ``arena`` + ``safe_zones`` + ``movement_features`` """ import json from pathlib import Path from typing import Any # ── Dataset source labels ────────────────────────────────────────────────── DATASET_SOURCE_OLD = "old" DATASET_SOURCE_SCOPE = "scavenger_hunt" # new schema, scavenger hunt DATASET_SOURCE_HIDE_SEEK = "hide_and_seek" # new schema, hide & seek DATASET_SOURCE_TAG = "tag" # new schema, tag def load_games_dataset(path: str) -> list[dict]: """Load the games dataset from JSON file. Args: path: Path to dataset JSON file Returns: List of game records """ with open(path, 'r', encoding='utf-8') as f: return json.load(f) # ── Schema detection helpers ──────────────────────────────────────────────── def _detect_schema(record: dict) -> str: """Detect which schema a record uses based on its keys. Detection priority: 1. ``expected_output`` key → old dataset 2. ``output`` key with ``hiding_zones`` → hide_and_seek 3. ``output`` key with ``arena`` → tag 4. ``output`` key with ``tasks`` → scavenger_hunt (new schema) 5. Fallback → old schema Returns one of ``DATASET_SOURCE_OLD``, ``DATASET_SOURCE_SCOPE``, ``DATASET_SOURCE_HIDE_SEEK``, or ``DATASET_SOURCE_TAG``. """ # Old dataset uses expected_output, new datasets use output if "expected_output" in record: return DATASET_SOURCE_OLD output = record.get("output", {}) if not output: return DATASET_SOURCE_OLD # Hide & seek has hiding_zones; tag has arena; scavenger hunt (new) has tasks if output.get("hiding_zones") is not None: return DATASET_SOURCE_HIDE_SEEK if output.get("arena") is not None: return DATASET_SOURCE_TAG if output.get("tasks") is not None: return DATASET_SOURCE_SCOPE return DATASET_SOURCE_OLD # ── Old schema normalizer ─────────────────────────────────────────────────── def _normalize_old_record(record: dict) -> dict: """Normalize a record from the old ``games_dataset.json`` schema. Expected structure: {input: {game_type, location: {city, area, location_type}, preferences: {duration_minutes, num_players, difficulty, age_group}}, expected_output: {rules: [str], tasks: [{task_id, description, location_hint, points, ...}], hints: [[str]], safety_flags: {is_safe, flags}}, metadata: {quality_score, source, notes}} """ inp = record.get("input", {}) loc = inp.get("location", {}) prefs = inp.get("preferences", {}) out = record.get("expected_output", {}) meta = record.get("metadata", {}) tasks = out.get("tasks", []) rules = out.get("rules", []) hints = out.get("hints", []) sf = out.get("safety_flags", {}) return { "id": record.get("id"), "dataset_source": DATASET_SOURCE_OLD, "game_type": inp.get("game_type", ""), "city": loc.get("city", ""), "area": loc.get("area", ""), "location_type": loc.get("location_type", ""), "duration_minutes": prefs.get("duration_minutes"), "num_players": prefs.get("num_players"), "difficulty": prefs.get("difficulty", ""), "age_group": prefs.get("age_group", ""), "landscape_tags": loc.get("landscape_tags", []), "theme": prefs.get("theme", ""), "mobility": prefs.get("mobility", "standard"), "allow_transport": prefs.get("allow_transport", False), "players_count": prefs.get("num_players"), "team_count": prefs.get("team_count", 1), "num_tasks": len(tasks), "task_ids": [t.get("task_id") for t in tasks], "num_rules": len(rules), "num_hints": len(hints), "is_safe": sf.get("is_safe", True), "safety_flags_list": sf.get("flags", []), "quality_score": meta.get("quality_score"), "source": meta.get("source"), "notes": meta.get("notes", ""), "rules": rules, "tasks": tasks, "hints": hints, # Hide & seek placeholders (not present in old schema) "hiding_zones": [], "play_area": {}, "seeker_guidance": {}, "hide_seek_scoring": {}, } # ── New schema normalizers ────────────────────────────────────────────────── def _normalize_new_scavenger_record(record: dict) -> dict: """Normalize a record from the new scavenger_hunt dataset. Expected structure: {input: {game_type, location: {city, country, city_code, landscape_tags, urban_density, climate_zone, area_type}, players: {count, team_count, age_group, mobility}, preferences: {duration_minutes, difficulty, theme, allow_transport}}, output: {game_title, rules: {objective, scoring_method, ...}, tasks: [{task_id, title, description, landscape_tags_used, task_type, difficulty_contribution, points, completion_proof, estimated_time_minutes, hints: {hint_1, hint_2, hint_3}, safety_flags}], safety_constraints: {...}, scoring_summary: {...}}} """ inp = record.get("input", {}) loc = inp.get("location", {}) players = inp.get("players", {}) prefs = inp.get("preferences", {}) out = record.get("output", {}) tasks = out.get("tasks", []) rules_obj = out.get("rules", {}) safety = out.get("safety_constraints", {}) # Flatten rules object into a list of rule strings rules_list = [] if isinstance(rules_obj, dict): if rules_obj.get("objective"): rules_list.append(f"Objective: {rules_obj['objective']}") if rules_obj.get("team_rules"): rules_list.append(rules_obj["team_rules"]) if rules_obj.get("disqualification_conditions"): for cond in rules_obj["disqualification_conditions"]: rules_list.append(f"No {cond}") elif isinstance(rules_obj, list): rules_list = rules_obj # Normalize each task: map new schema fields to unified keys normalized_tasks = [] for t in tasks: nt = dict(t) # Map completion_proof → proof_type (new SH schema) if "completion_proof" in nt and "proof_type" not in nt: nt["proof_type"] = nt["completion_proof"] normalized_tasks.append(nt) return { "id": record.get("id"), "dataset_source": DATASET_SOURCE_SCOPE, "game_type": inp.get("game_type", "scavenger_hunt"), "city": loc.get("city", ""), "area": loc.get("area_type", loc.get("city", "")), "location_type": loc.get("area_type", ""), "duration_minutes": prefs.get("duration_minutes"), "num_players": players.get("count"), "difficulty": prefs.get("difficulty", ""), "age_group": players.get("age_group", ""), "landscape_tags": loc.get("landscape_tags", []), "theme": prefs.get("theme", ""), "mobility": players.get("mobility", "standard"), "allow_transport": prefs.get("allow_transport", False), "players_count": players.get("count"), "team_count": players.get("team_count", 1), "num_tasks": len(tasks), "task_ids": [t.get("task_id") for t in tasks], "num_rules": len(rules_list), "num_hints": sum(1 for t in tasks if t.get("hints")), "is_safe": True, "safety_flags_list": safety.get("exclusion_zones", []) + safety.get("physical_limits", []), "quality_score": out.get("quality_score"), "source": "new_scavenger", "notes": safety.get("notes", ""), "rules": rules_list, "tasks": normalized_tasks, # normalized tasks with unified fields "hints": [t.get("hints", {}) for t in tasks], # Hide & seek placeholders "hiding_zones": [], "play_area": {}, "seeker_guidance": {}, "hide_seek_scoring": {}, } def _normalize_new_hide_seek_record(record: dict) -> dict: """Normalize a record from the new hide_and_seek dataset. Expected structure: {input: {game_type, location: {city, country, city_code, landscape_tags, ...}, players: {count, team_count, age_group, mobility}, preferences: {duration_minutes, difficulty, theme, allow_transport}}, output: {game_title, rules: {objective, round_structure, round_count, hide_time_seconds, round_time_minutes, seeker_count, win_condition}, play_area: {boundary_description, boundary_landscape_tags_used, safe_base_description, boundary_size_tier, out_of_bounds}, hiding_zones: [{zone_id, description, landscape_tags_used, concealment_rating, risk_note, safety_flags}], seeker_guidance: {search_strategy, high_probability_zones}, safety_constraints: {...}, scoring: {...}, ...}} """ inp = record.get("input", {}) loc = inp.get("location", {}) players = inp.get("players", {}) prefs = inp.get("preferences", {}) out = record.get("output", {}) rules_obj = out.get("rules", {}) hiding_zones = out.get("hiding_zones", []) safety = out.get("safety_constraints", {}) # Flatten rules rules_list = [] if isinstance(rules_obj, dict): if rules_obj.get("objective"): rules_list.append(f"Objective: {rules_obj['objective']}") if rules_obj.get("round_structure"): rules_list.append(rules_obj["round_structure"]) if rules_obj.get("win_condition"): rules_list.append(f"Win: {rules_obj['win_condition']}") elif isinstance(rules_obj, list): rules_list = rules_obj return { "id": record.get("id"), "dataset_source": DATASET_SOURCE_HIDE_SEEK, "game_type": inp.get("game_type", "hide_and_seek"), "city": loc.get("city", ""), "area": loc.get("area_type", loc.get("city", "")), "location_type": loc.get("area_type", ""), "duration_minutes": prefs.get("duration_minutes"), "num_players": players.get("count"), "difficulty": prefs.get("difficulty", ""), "age_group": players.get("age_group", ""), "landscape_tags": loc.get("landscape_tags", []), "theme": prefs.get("theme", ""), "mobility": players.get("mobility", "standard"), "allow_transport": prefs.get("allow_transport", False), "players_count": players.get("count"), "team_count": players.get("team_count", 1), "num_tasks": len(hiding_zones), "task_ids": [z.get("zone_id") for z in hiding_zones], "num_rules": len(rules_list), "num_hints": 0, "is_safe": True, "safety_flags_list": safety.get("exclusion_zones", []) + safety.get("physical_limits", []), "quality_score": out.get("quality_score"), "source": "new_hide_seek", "notes": safety.get("notes", ""), "rules": rules_list, "tasks": hiding_zones, # hiding_zones stored as 'tasks' for unified retrieval "hints": [], # Hide & seek specific fields "hiding_zones": hiding_zones, "play_area": out.get("play_area", {}), "seeker_guidance": out.get("seeker_guidance", {}), "hide_seek_scoring": out.get("scoring", {}), # Tag placeholders "arena": {}, "safe_zones": [], "movement_features": {}, } # ── New tag normalizer ──────────────────────────────────────────────────── def _normalize_new_tag_record(record: dict) -> dict: """Normalize a record from the new tag dataset. Expected structure: {input: {game_type, location: {city, ..., landscape_tags, ...}, players: {count, team_count, age_group, mobility}, preferences: {duration_minutes, difficulty, theme, allow_transport}}, output: {game_title, rules: {objective, tag_variant, it_count, round_count, ...}, arena: {boundary_description, boundary_size_tier, out_of_bounds}, safe_zones: [{zone_id, description, capacity_rule, safety_flags}], movement_features: {chokepoints, open_zones, hazard_zones}, safety_constraints: {...}, scoring: {...}, ...}} """ inp = record.get("input", {}) loc = inp.get("location", {}) players = inp.get("players", {}) prefs = inp.get("preferences", {}) out = record.get("output", {}) rules_obj = out.get("rules", {}) safety = out.get("safety_constraints", {}) arena = out.get("arena", {}) safe_zones = out.get("safe_zones", []) # Flatten rules rules_list = [] if isinstance(rules_obj, dict): if rules_obj.get("objective"): rules_list.append(f"Objective: {rules_obj['objective']}") if rules_obj.get("tag_mechanic"): rules_list.append(f"Mechanic: {rules_obj['tag_mechanic']}") if rules_obj.get("unfreeze_or_revival_rule"): rules_list.append(f"Revival: {rules_obj['unfreeze_or_revival_rule']}") if rules_obj.get("win_condition"): rules_list.append(f"Win: {rules_obj['win_condition']}") elif isinstance(rules_obj, list): rules_list = rules_obj return { "id": record.get("id"), "dataset_source": DATASET_SOURCE_TAG, "game_type": inp.get("game_type", "tag"), "city": loc.get("city", ""), "area": loc.get("area_type", loc.get("city", "")), "location_type": loc.get("area_type", ""), "duration_minutes": prefs.get("duration_minutes"), "num_players": players.get("count"), "difficulty": prefs.get("difficulty", ""), "age_group": players.get("age_group", ""), "landscape_tags": loc.get("landscape_tags", []), "theme": prefs.get("theme", ""), "mobility": players.get("mobility", "standard"), "allow_transport": prefs.get("allow_transport", False), "players_count": players.get("count"), "team_count": players.get("team_count", 1), "num_tasks": len(safe_zones), "task_ids": [z.get("zone_id") for z in safe_zones], "num_rules": len(rules_list), "num_hints": 0, "is_safe": True, "safety_flags_list": safety.get("exclusion_zones", []) + safety.get("physical_limits", []), "quality_score": out.get("quality_score"), "source": "new_tag", "notes": safety.get("notes", ""), "rules": rules_list, "tasks": safe_zones, "hints": [], # Tag specific fields "arena": arena, "safe_zones": safe_zones, "movement_features": out.get("movement_features", {}), "tag_variant": rules_obj.get("tag_variant", "classic"), "it_count": rules_obj.get("it_count", 1), "round_count": rules_obj.get("round_count", 1), "tag_mechanic": rules_obj.get("tag_mechanic", ""), "tag_scoring": out.get("scoring", {}), # Placeholders for other schemas "hiding_zones": [], "play_area": {}, "seeker_guidance": {}, "hide_seek_scoring": {}, } # ── Dispatcher ────────────────────────────────────────────────────────────── def normalize_game_record(record: dict) -> dict: """Normalize a raw game record into structured format. Auto-detects the schema (old, new scavenger hunt, new hide & seek, or new tag) and dispatches to the appropriate normalizer. Args: record: Raw game record from any dataset Returns: Normalized game record with unified fields """ schema = _detect_schema(record) if schema == DATASET_SOURCE_TAG: return _normalize_new_tag_record(record) if schema == DATASET_SOURCE_HIDE_SEEK: return _normalize_new_hide_seek_record(record) if schema == DATASET_SOURCE_SCOPE: return _normalize_new_scavenger_record(record) return _normalize_old_record(record) # ── Similarity scoring ───────────────────────────────────────────────────── def _compute_similarity_score(config: dict, record: dict) -> float: """Compute similarity score between a user config and a dataset record. Awards bonus points for: - Same game type (highest weight, +50) - New dataset source preferred (+15) - Age group match (+25/+15) - Difficulty match (+20) - Location type match (+15) - Landscape tag overlap (+10) - City name match (+10) - Duration proximity (+10/+5) - Theme match (+8) - Mobility match (+5) - Transport preference match (+5) - Quality bonus (+0.5 × score) - Safety bonus (+5) """ score = 0.0 # ── Game type exact match (highest weight) ─────────────────────── if config.get('game_type', '').lower() == record.get('game_type', '').lower(): score += 50 # ── Dataset source bonus — prefer new, richer datasets ────────── src = record.get('dataset_source', DATASET_SOURCE_OLD) if src in (DATASET_SOURCE_SCOPE, DATASET_SOURCE_HIDE_SEEK, DATASET_SOURCE_TAG): score += 15 # ── City match ────────────────────────────────────────────────── config_city = config.get('city', '').lower().strip() record_city = record.get('city', '').lower().strip() if config_city and record_city: if config_city == record_city: score += 10 # same city is a strong signal # ── Landscape tags overlap ────────────────────────────────────── config_tags = set(config.get('landscape_tags', [])) record_tags = set(record.get('landscape_tags', [])) if config_tags and record_tags: overlap = len(config_tags & record_tags) if overlap > 0: score += min(overlap * 3, 10) # max +10 # ── Age group match ───────────────────────────────────────────── config_age = config.get('age_group', '').lower() record_age = record.get('age_group', '').lower() if config_age and record_age: if config_age == record_age: score += 25 elif 'mixed' in [config_age, record_age]: score += 15 # ── Difficulty match ──────────────────────────────────────────── if config.get('difficulty', '').lower() == record.get('difficulty', '').lower(): score += 20 # ── Location type match ───────────────────────────────────────── if config.get('location_type', '').lower() == record.get('location_type', '').lower(): score += 15 # ── Theme match ───────────────────────────────────────────────── config_theme = config.get('theme', '').lower().strip() record_theme = record.get('theme', '').lower().strip() if config_theme and record_theme and config_theme == record_theme: score += 8 # ── Duration proximity ────────────────────────────────────────── config_duration = config.get('duration_minutes') record_duration = record.get('duration_minutes') if config_duration and record_duration: duration_diff = abs(config_duration - record_duration) if duration_diff <= 15: score += 10 elif duration_diff <= 30: score += 5 # ── Mobility match ────────────────────────────────────────────── config_mob = config.get('mobility', '').lower() record_mob = record.get('mobility', '').lower() if config_mob and record_mob and config_mob == record_mob: score += 5 # ── Transport preference match ────────────────────────────────── config_transport = config.get('allow_transport') record_transport = record.get('allow_transport') if config_transport is not None and record_transport is not None: if config_transport == record_transport: score += 5 # ── Quality bonus ─────────────────────────────────────────────── quality = record.get('quality_score', 0) if quality: score += quality * 0.5 # ── Safety bonus ──────────────────────────────────────────────── if record.get('is_safe', True): score += 5 return score # ── Game-type-aware exemplar builder ─────────────────────────────────────── def _build_exemplar(record: dict, retrieval_score: float) -> dict: """Build a compressed exemplar bundle from a normalized record. The exemplar shape differs by game type so the generation prompt receives relevant context. """ base = { 'id': record['id'], 'dataset_source': record.get('dataset_source', DATASET_SOURCE_OLD), 'game_type': record['game_type'], 'area': record.get('area', ''), 'city': record.get('city', ''), 'difficulty': record.get('difficulty', ''), 'age_group': record.get('age_group', ''), 'duration_minutes': record.get('duration_minutes'), 'location_type': record.get('location_type', ''), 'landscape_tags': record.get('landscape_tags', []), 'theme': record.get('theme', ''), 'mobility': record.get('mobility', 'standard'), 'num_players': record.get('num_players') or record.get('players_count'), 'rules_summary': record.get('rules', [])[:3], 'safety_patterns': record.get('safety_flags_list', []), 'notes': record.get('notes', '')[:120], 'quality_score': record.get('quality_score'), 'retrieval_score': retrieval_score, } game_type = record.get('game_type', '') if game_type == 'hide_and_seek': # Hide & seek: expose hiding zones, play area, seeker guidance base['hiding_zones_summary'] = [ { 'zone_id': z.get('zone_id'), 'description': z.get('description', '')[:120], 'concealment_rating': z.get('concealment_rating', 'medium'), 'landscape_tags_used': z.get('landscape_tags_used', []), } for z in record.get('hiding_zones', [])[:3] ] base['play_area_summary'] = { k: record.get('play_area', {}).get(k, '') for k in ('boundary_description', 'boundary_size_tier', 'out_of_bounds') if record.get('play_area', {}).get(k) } base['seeker_strategy'] = ( record.get('seeker_guidance', {}).get('search_strategy', '')[:200] ) base['hide_seek_scoring'] = record.get('hide_seek_scoring', {}) base['task_patterns'] = [] elif game_type == 'tag' and record.get('dataset_source') == DATASET_SOURCE_TAG: # Tag (new schema): expose arena, safe zones, movement features arena = record.get('arena', {}) base['arena_summary'] = { k: arena.get(k, '') for k in ('boundary_description', 'arena_size_tier', 'out_of_bounds') if arena.get(k) } base['safe_zones_summary'] = [ { 'zone_id': z.get('zone_id'), 'description': z.get('description', '')[:120], 'capacity_rule': z.get('capacity_rule', ''), } for z in record.get('safe_zones', [])[:3] ] mf = record.get('movement_features', {}) base['chokepoints'] = mf.get('chokepoints', [])[:3] base['open_zones'] = mf.get('open_zones', [])[:3] base['tag_variant'] = record.get('tag_variant', 'classic') base['it_count'] = record.get('it_count', 1) base['round_count'] = record.get('round_count', 1) base['tag_mechanic'] = record.get('tag_mechanic', '')[:150] base['tag_scoring'] = record.get('tag_scoring', {}) base['task_patterns'] = [] else: # Scavenger hunt (old or new): expose task patterns tasks = record.get('tasks', []) base['task_patterns'] = [ { 'task_id': t.get('task_id'), 'title': t.get('title', ''), 'difficulty': t.get('difficulty_contribution', ''), 'task_type': t.get('task_type', ''), 'points': t.get('points'), 'proof_type': t.get('proof_type', 'observation'), 'time_limit': t.get('estimated_time_minutes') or t.get('time_limit_minutes'), 'landscape_tags_used': t.get('landscape_tags_used', []), } for t in tasks[:3] ] # Ensure all exemplars have these keys (even if empty) base.setdefault('hiding_zones_summary', []) base.setdefault('task_patterns', []) return base # ── Public retrieval API ──────────────────────────────────────────────────── def retrieve_examples(config: dict, dataset: list[dict], k: int = 3) -> list[dict]: """Retrieve top k closest examples from dataset. Uses a multi-factor similarity score (game type, city, tags, age group, difficulty, duration, theme, mobility) to rank all records, then returns game-type-aware compressed exemplar bundles. Args: config: Game configuration from user input dataset: Loaded and normalized dataset records (from all sources) k: Number of examples to retrieve Returns: List of retrieved example games (compressed exemplar bundles) """ if not dataset: return [] requested_type = config.get('game_type', '').lower() # Score all records scored_records = [] for record in dataset: # Primary filter: must match game type (same-type examples are most useful) record_type = record.get('game_type', '').lower() if record_type != requested_type: continue score = _compute_similarity_score(config, record) scored_records.append({'record': record, 'score': score}) # Sort by score descending scored_records.sort(key=lambda x: x['score'], reverse=True) # Take top k and build exemplar bundles top_k = scored_records[:k] return [_build_exemplar(item['record'], item['score']) for item in top_k]