feat: enable real case website e2e capture

scripts/capture_commercial_astrology_e2e_contexts.py

@@ -68,14 +68,57 @@ def _capture_body(question: dict[str, Any]) -> dict[str, Any]:
     }
 
 
-def capture(contract_path: Path = DEFAULT_CONTRACT, output_dir: Path = DEFAULT_OUTPUT_DIR) -> dict[str, Any]:
+def _capture_body_for_real_case(case: dict[str, Any], prompt: str) -> dict[str, Any]:
+    birth = case["birth"]
+    year, month, day = [int(part) for part in birth["date"].split("-")]
+    hour, minute = [int(part) for part in birth["time"].split(":")[:2]]
+    theme_map = {
+        "timing": "career",
+        "annual": "career",
+        "migration": "career",
+        "family": "marriage",
+        "education": "career",
+    }
+    themes = [theme_map.get(domain, domain) for domain in case["domains"]]
+    return {
+        "year": year,
+        "month": month,
+        "day": day,
+        "hour": hour,
+        "minute": minute,
+        "lat": birth["lat"],
+        "lon": birth["lon"],
+        "tz": birth["tz"],
+        "city": birth["place"],
+        "question": prompt,
+        "question_text": prompt,
+        "theme": themes,
+        "evaluation_domains": case["domains"],
+        "entry_mode": "direct_chart",
+        "case_id": case["case_id"],
+        "subject": case["subject"],
+        "source_policy": birth["source_policy"],
+    }
+
+
+def capture(contract_path: Path = DEFAULT_CONTRACT, output_dir: Path = DEFAULT_OUTPUT_DIR, max_items: int | None = None) -> dict[str, Any]:
     contract = _load_json(contract_path)
     output_dir = output_dir.resolve()
     output_dir.mkdir(parents=True, exist_ok=True)
     rows: list[dict[str, Any]] = []
-    for question in contract["questions"]:
+    if "cases" in contract:
+        questions = [
+            {"id": f"{case['case_id']}__{index + 1}", "body": _capture_body_for_real_case(case, prompt)}
+            for case in contract["cases"]
+            for index, prompt in enumerate(case["prompts"])
+        ]
+    else:
+        questions = [{"id": str(question["id"]), "body": _capture_body(question)} for question in contract["questions"]]
+    for question in questions:
+        if max_items is not None and len(rows) >= max_items:
+            break
         qid = str(question["id"])
-        result = execute_consultation_workflow(_capture_body(question), surface="commercial_e2e_capture")
+        result = execute_consultation_workflow(question["body"], surface="commercial_e2e_capture")
         context_path = output_dir / f"{qid}.json"
         context_path.write_text(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
         rows.append(
@@ -106,8 +149,9 @@ def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--contract", type=Path, default=DEFAULT_CONTRACT)
     parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
+    parser.add_argument("--max-items", type=int, default=None)
     args = parser.parse_args()
-    print(json.dumps(capture(args.contract, args.output_dir), ensure_ascii=False, indent=2, sort_keys=True))
+    print(json.dumps(capture(args.contract, args.output_dir, max_items=args.max_items), ensure_ascii=False, indent=2, sort_keys=True))
     return 0