Update usage_logging.py

usage_logging.py

@@ -145,32 +145,35 @@ def _country_lookup(ip: str) -> tuple[str, str]:
         return ("", "")
 
 
-def append_visit_to_dataset(country: str, city: str, event_type: str = "usage_start", **extra_fields):
+def append_visit_to_dataset(
+    country: str,
+    city: str,
+    event_type: str = "usage_start",
+    country_source: str = "unknown",
+    country_code: str = "",
+    **extra_fields
+):
     api = _hf_api()
     if not api:
         return
 
-    existing = _download_text_if_exists(
-        USAGE_DATASET_REPO,
-        USAGE_JSONL_PATH
-    )
+    existing = _download_text_if_exists(USAGE_DATASET_REPO, USAGE_JSONL_PATH)
 
     event = {
         "ts_utc": datetime.utcnow().isoformat() + "Z",
         "space_url": SPACE_URL,
-        "country": country,
-        "city": city,
+        "country": country or "Unknown",
+        "country_code": (country_code or "").strip().upper(),
+        "country_source": country_source or "unknown",
+        "city": city or "",
         "event": event_type,
     }
 
+    # Drop None values so you never write JSON nulls
     if extra_fields:
-        event.update(extra_fields)
+        event.update({k: v for k, v in extra_fields.items() if v is not None})
 
-    new_content = (
-        existing.rstrip("\n") + "\n"
-        if existing.strip()
-        else ""
-    ) + json.dumps(event) + "\n"
+    new_content = (existing.rstrip("\n") + "\n" if existing.strip() else "") + json.dumps(event) + "\n"
 
     try:
         api.upload_file(
@@ -189,11 +192,11 @@ def record_visit(request: gr.Request):
     country_hint = _country_from_headers(request)
     if _is_valid_country_code(country_hint):
         append_visit_to_dataset(
-            country=_expand_country_code(country_hint),
-            city="",
-            event_type="usage_start",
-            country_source="header",
-            country_code=country_hint.strip().upper(),
+        country=_expand_country_code(country_hint),
+        city="",
+        event_type="usage_start",
+        country_source="header",
+        country_code=country_hint,
         )
         return
 
@@ -203,11 +206,11 @@ def record_visit(request: gr.Request):
         cc, name = _country_lookup(ip)
         if _is_valid_country_code(cc):
             append_visit_to_dataset(
-                country=name or _expand_country_code(cc),
-                city="",
-                event_type="usage_start",
-                country_source="ipinfo",
-                country_code=cc,
+            country=name or _expand_country_code(cc),
+            city="",
+            event_type="usage_start",
+            country_source="ipinfo",
+            country_code=cc,
             )
         else:
             append_visit_to_dataset(