Update usage_logging.py

usage_logging.py

@@ -276,12 +276,16 @@ def _expand_country_code(code: str) -> str:
 def rebuild_visits_rollup_from_event_files() -> str:
     """
     Rebuilds usage/visits.jsonl from immutable per-event JSON files in usage/events/.
-    This is safe if triggered manually (admin button).
+    ALSO writes an enriched rollup usage/visits_enriched.jsonl where:
+      - legacy rows (no session_id) keep their original country as final_country
+      - new click rows (with session_id) get final_country from the session's usage_start row when available
     """
     api = _hf_api()
     if not api:
         return "HF_TOKEN not available. Rollup requires write access."
 
+    ENRICHED_ROLLUP_PATH = "usage/visits_enriched.jsonl"
+
     # 1) List files
     try:
         files = list_repo_files(repo_id=USAGE_DATASET_REPO, repo_type="dataset")
@@ -317,22 +321,91 @@ def rebuild_visits_rollup_from_event_files() -> str:
     # 3) Sort by ts_utc
     events.sort(key=lambda e: (e.get("ts_utc") or ""))
 
-    # 4) Write JSONL
-    buf = BytesIO()
+    # 4) Build session → geo map from usage_start events (new schema)
+    # NOTE: you said you are using "usage_start" (not "session_start")
+    session_geo = {}
+    for e in events:
+        if e.get("event") == "usage_start":
+            sid = (e.get("session_id") or "").strip()
+            if not sid:
+                continue  # legacy usage_start rows won’t have session_id; ignore for mapping
+            if sid not in session_geo:
+                session_geo[sid] = {
+                    "country_session": e.get("country") or "Unknown",
+                    "country_code_session": (e.get("country_code") or "").strip().upper(),
+                    "country_source_session": e.get("country_source") or "unknown",
+                }
+
+    # 5) Write RAW JSONL (same behavior as today)
+    buf_raw = BytesIO()
+    for evt in events:
+        buf_raw.write((json.dumps(evt, ensure_ascii=False) + "\n").encode("utf-8"))
+    buf_raw.seek(0)
+
+    # 6) Write ENRICHED JSONL
+    # Rules:
+    # - Legacy rows (no session_id): final_country = evt.country (already correct)
+    # - New rows with session_id:
+    #     - if evt.country is Unknown/blank, use country from session usage_start
+    #     - keep original evt.country as well (don’t overwrite)
+    buf_enriched = BytesIO()
+
     for evt in events:
-        buf.write((json.dumps(evt, ensure_ascii=False) + "\n").encode("utf-8"))
-    buf.seek(0)
+        enriched = dict(evt)
+        sid = (evt.get("session_id") or "").strip()
+        geo = session_geo.get(sid, {}) if sid else {}
+
+        # Keep a copy of session geo (useful for debugging/audit)
+        if geo:
+            enriched.update(geo)
+
+        evt_country = (evt.get("country") or "").strip()
+        evt_cc = (evt.get("country_code") or "").strip().upper()
+
+        # Determine final_country / final_country_code
+        if sid:
+            # New schema row: prefer existing event country if it’s meaningful,
+            # else fall back to session geo.
+            if evt_country and evt_country.lower() != "unknown":
+                enriched["final_country"] = evt_country
+                enriched["final_country_code"] = evt_cc
+                enriched["final_country_source"] = evt.get("country_source") or "unknown"
+            else:
+                enriched["final_country"] = geo.get("country_session", "Unknown")
+                enriched["final_country_code"] = geo.get("country_code_session", "")
+                enriched["final_country_source"] = geo.get("country_source_session", "unknown")
+        else:
+            # Legacy row: preserve the original country fields
+            enriched["final_country"] = evt_country or "Unknown"
+            enriched["final_country_code"] = evt_cc
+            enriched["final_country_source"] = evt.get("country_source") or "unknown"
+
+        buf_enriched.write((json.dumps(enriched, ensure_ascii=False) + "\n").encode("utf-8"))
+
+    buf_enriched.seek(0)
 
-    # 5) Upload rollup
+    # 7) Upload both rollups
     try:
         api.upload_file(
             repo_id=USAGE_DATASET_REPO,
             repo_type="dataset",
-            path_in_repo=ROLLUP_PATH,
-            path_or_fileobj=buf,
+            path_in_repo=ROLLUP_PATH,  # your existing usage/visits.jsonl
+            path_or_fileobj=buf_raw,
             commit_message=f"rebuild {ROLLUP_PATH} from {USAGE_EVENTS_DIR}",
         )
+        api.upload_file(
+            repo_id=USAGE_DATASET_REPO,
+            repo_type="dataset",
+            path_in_repo=ENRICHED_ROLLUP_PATH,
+            path_or_fileobj=buf_enriched,
+            commit_message=f"rebuild {ENRICHED_ROLLUP_PATH} from {USAGE_EVENTS_DIR}",
+        )
     except Exception as e:
         return f"Rollup upload failed: {repr(e)}"
 
-    return f"Rollup rebuilt: {ROLLUP_PATH} rows={len(events)} (bad_files={bad})."
+    return (
+        f"Rollups rebuilt:\n"
+        f"- RAW: {ROLLUP_PATH} rows={len(events)} (bad_files={bad})\n"
+        f"- ENRICHED: {ENRICHED_ROLLUP_PATH} rows={len(events)} (bad_files={bad})"
+    )
+