from __future__ import annotations import json from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Literal DecisionLabel = Literal["do_now", "do_later", "remind", "no_action"] MissingFieldStrategy = Literal["unknown", "synthetic"] _USER_STATE_FLAG_TO_TOKEN = { "user_asleep": "asleep", "user_in_rush": "in_rush", "user_injured_or_disabled": "injured_or_disabled", "user_nearby": "nearby", } _ENV_FLAG_DIRECT = { "adverse_weather": "adverse_weather", "guests_present": "guests_present", } _VALID_DECISIONS = {"do_now", "do_later", "remind", "no_action"} @dataclass class BuilderConfig: language: Literal["en", "es"] = "en" missing_field_strategy: MissingFieldStrategy = "unknown" include_compact_task_fields: bool = True @dataclass class BuiltTrainingSample: sample_id: str user_id: int user_external_id: str label_action: DecisionLabel action_input: dict[str, Any] context_input: dict[str, Any] structured_task_features: dict[str, Any] preference_snapshot: list[dict[str, Any]] = field(default_factory=list) source_metadata: dict[str, Any] = field(default_factory=dict) data_provenance: dict[str, str] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: return asdict(self) class TrainingDataBuilder: """ Builds model-ready samples from mapped TAACO JSONL files. Designed for MVP bootstrap training where feedback events are still sparse. """ def __init__(self, config: BuilderConfig | None = None): self.config = config or BuilderConfig() def build_from_taaco_mapped( self, *, input_path: Path, ) -> list[BuiltTrainingSample]: records = self._load_jsonl(input_path) persona_map = self._build_persona_id_map(records) out: list[BuiltTrainingSample] = [] for row in records: sample = self._build_sample(row=row, persona_id_map=persona_map) if sample is not None: out.append(sample) return out def write_jsonl(self, samples: list[BuiltTrainingSample], output_path: Path) -> None: output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8") as f: for sample in samples: f.write(json.dumps(sample.to_dict(), ensure_ascii=False) + "\n") def summarize(self, samples: list[BuiltTrainingSample]) -> dict[str, Any]: label_counts: dict[str, int] = {} user_counts: dict[str, int] = {} with_time = 0 with_weekday = 0 with_available_objects = 0 for sample in samples: label_counts[sample.label_action] = label_counts.get(sample.label_action, 0) + 1 user_counts[sample.user_external_id] = ( user_counts.get(sample.user_external_id, 0) + 1 ) if sample.context_input.get("time_of_day", "unknown") != "unknown": with_time += 1 if sample.context_input.get("weekday", "unknown") != "unknown": with_weekday += 1 if sample.context_input.get("available_objects"): with_available_objects += 1 total = len(samples) return { "num_samples": total, "labels": label_counts, "users": user_counts, "coverage": { "time_of_day_known": with_time, "weekday_known": with_weekday, "available_objects_nonempty": with_available_objects, }, "config": asdict(self.config), } def _build_sample( self, *, row: dict[str, Any], persona_id_map: dict[str, int], ) -> BuiltTrainingSample | None: persona = str(row.get("source_persona", "unknown_persona")) user_id = persona_id_map[persona] task_input = self._select_task_input(row) if task_input is None: return None label = self._select_label(row) if label is None: return None action_text = str(task_input.get("action", "")).strip() activity = str(task_input.get("activity", "")).strip() or None locations = self._as_list(task_input.get("locations")) objects = self._as_list(task_input.get("objects")) conditions = self._as_list(task_input.get("conditions")) context_flags = task_input.get("context_flags") or {} context_input, provenance = self._build_context_input( locations=locations, objects=objects, conditions=conditions, context_flags=context_flags, user_busy=bool(task_input.get("user_busy", False)), quiet_hours=bool(task_input.get("quiet_hours", False)), ) sample_id = f"{persona}:{task_input.get('task_id', '')}:{row.get('source_preference_index', '')}" structured_task = { "kind": task_input.get("kind"), "urgency": task_input.get("urgency"), "sensitivity": task_input.get("sensitivity"), "user_busy": bool(task_input.get("user_busy", False)), "quiet_hours": bool(task_input.get("quiet_hours", False)), } if self.config.include_compact_task_fields: structured_task["conditions"] = conditions structured_task["context_flags"] = context_flags return BuiltTrainingSample( sample_id=sample_id, user_id=user_id, user_external_id=persona, label_action=label, action_input={ "action_text": action_text, "activity": activity, }, context_input=context_input, structured_task_features=structured_task, preference_snapshot=[], source_metadata={ "source_example_index": row.get("source_example_index"), "source_preference_index": row.get("source_preference_index"), "expected_action_taaco": self._extract_expected_action_taaco(row), "mapping_trace": row.get("mapping_trace", []), }, data_provenance=provenance, ) def _build_context_input( self, *, locations: list[str], objects: list[str], conditions: list[str], context_flags: dict[str, Any], user_busy: bool, quiet_hours: bool, ) -> tuple[dict[str, Any], dict[str, str]]: true_flags = {k for k, v in context_flags.items() if bool(v)} user_state: list[str] = [] for flag, token in _USER_STATE_FLAG_TO_TOKEN.items(): if flag in true_flags: user_state.append(token) if user_busy: user_state.append("busy") environment_flags: list[str] = [] for flag, token in _ENV_FLAG_DIRECT.items(): if flag in true_flags: environment_flags.append(token) if quiet_hours: environment_flags.append("quiet_hours") if "weekend" in true_flags: environment_flags.append("weekend") location_current = locations[0] if locations else None time_of_day = "unknown" weekday = "unknown" available_objects: list[str] = [] provenance = { "location_current": "derived_from_locations[0]", "objects_nearby": "direct_from_task.objects", "raw_conditions": "direct_from_task.conditions", "user_state": "derived_from_context_flags_and_user_busy", "environment_flags": "derived_from_context_flags_and_quiet_hours", } if self.config.missing_field_strategy == "synthetic": available_objects = sorted(set(objects)) provenance["available_objects"] = "synthetic_from_objects_nearby" if "early_morning" in true_flags: time_of_day = "morning" provenance["time_of_day"] = "synthetic_from_context_flags.early_morning" elif "user_asleep" in true_flags: time_of_day = "night" provenance["time_of_day"] = "synthetic_from_context_flags.user_asleep" else: provenance["time_of_day"] = "unknown_no_signal" if "weekend" in true_flags: weekday = "saturday" provenance["weekday"] = "synthetic_from_context_flags.weekend" else: provenance["weekday"] = "unknown_no_signal" else: provenance["available_objects"] = "unknown_default" provenance["time_of_day"] = "unknown_default" provenance["weekday"] = "unknown_default" return ( { "location_current": location_current, "objects_nearby": objects, "available_objects": available_objects, "raw_conditions": conditions, "time_of_day": time_of_day, "weekday": weekday, "user_state": sorted(set(user_state)), "environment_flags": sorted(set(environment_flags)), }, provenance, ) def _select_task_input(self, row: dict[str, Any]) -> dict[str, Any] | None: # Preferred format: all_personas_mapped_both.jsonl if self.config.language == "en" and isinstance(row.get("task_input_en"), dict): return row["task_input_en"] if self.config.language == "es" and isinstance(row.get("task_input_es"), dict): return row["task_input_es"] # Compact format: all_personas_mapped_en/es.jsonl compact = row.get("task_input") source = row.get("source") or {} if isinstance(compact, dict): if "action" not in compact and isinstance(source, dict): # Recover dropped fields from source payload. compact = dict(compact) compact["action"] = source.get("action", "") compact["activity"] = source.get("activity", "") compact["objects"] = self._as_list(source.get("object")) compact["locations"] = self._as_list(source.get("location")) compact["conditions"] = self._as_list(source.get("conditions")) compact["explanations"] = source.get("explanations", []) compact["explanations_opposing"] = source.get("explanations_opposing", []) compact["context_flags"] = row.get("context_flags", {}) return compact return None def _select_label(self, row: dict[str, Any]) -> DecisionLabel | None: expected = row.get("expected_action") if isinstance(expected, str): return self._normalize_label(expected) if isinstance(expected, dict): if self.config.language == "en": return self._normalize_label(expected.get("en")) if self.config.language == "es": # Training label remains canonical English token. return self._normalize_label(expected.get("taaco")) return None def _extract_expected_action_taaco(self, row: dict[str, Any]) -> str | None: expected = row.get("expected_action") if isinstance(expected, str): return expected if isinstance(expected, dict): taaco = expected.get("taaco") return str(taaco) if taaco is not None else None alt = row.get("expected_action_taaco") return str(alt) if alt is not None else None def _build_persona_id_map(self, rows: list[dict[str, Any]]) -> dict[str, int]: personas = sorted({str(r.get("source_persona", "unknown_persona")) for r in rows}) return {persona: i + 1 for i, persona in enumerate(personas)} def _normalize_label(self, label: Any) -> DecisionLabel | None: if label is None: return None text = str(label).strip().lower().replace(" ", "_") if text in _VALID_DECISIONS: return text # type: ignore[return-value] return None def _as_list(self, value: Any) -> list[str]: if value is None: return [] if isinstance(value, list): return [str(x) for x in value] return [str(value)] def _load_jsonl(self, path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as f: for raw in f: line = raw.strip() if not line: continue rows.append(json.loads(line)) return rows __all__ = [ "BuilderConfig", "BuiltTrainingSample", "DecisionLabel", "MissingFieldStrategy", "TrainingDataBuilder", ]