Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import os | |
| import sqlite3 | |
| import uuid | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any, Callable | |
| from hf_dataset_logs import hf_log_status, upload_run_log_to_hf_dataset | |
| APP_DIR = Path(__file__).resolve().parent | |
| def default_log_path() -> Path: | |
| configured = os.getenv("SILICON_LOG_DB", "").strip() | |
| if configured: | |
| return Path(configured).expanduser() | |
| data_dir = Path("/data") | |
| if data_dir.exists() and os.access(data_dir, os.W_OK): | |
| return data_dir / "silicon_logs.sqlite3" | |
| return APP_DIR / "runs" / "silicon_logs.sqlite3" | |
| def log_status() -> dict[str, Any]: | |
| path = default_log_path() | |
| return { | |
| "enabled": True, | |
| "path": str(path), | |
| "exists": path.exists(), | |
| "storage": "persistent_bucket" if str(path).startswith("/data/") else "local_ephemeral", | |
| "runCount": count_runs(path) if path.exists() else 0, | |
| "hfDataset": hf_log_status(), | |
| } | |
| def record_run( | |
| payload: dict[str, Any], | |
| result: dict[str, Any], | |
| client: dict[str, Any] | None = None, | |
| hf_uploader: Callable[[dict[str, Any]], dict[str, Any]] | None = None, | |
| ) -> dict[str, Any]: | |
| path = default_log_path() | |
| ensure_schema(path) | |
| run_id = f"run_{uuid.uuid4().hex[:16]}" | |
| created_at = datetime.now(timezone.utc).isoformat() | |
| config = payload.get("config", {}) if isinstance(payload.get("config"), dict) else {} | |
| run_document = build_run_document( | |
| run_id=run_id, | |
| created_at=created_at, | |
| config=config, | |
| result=result, | |
| client=client or {}, | |
| ) | |
| with sqlite3.connect(path) as conn: | |
| conn.execute( | |
| """ | |
| INSERT INTO runs( | |
| run_id, | |
| created_at, | |
| sample_size, | |
| respondent_settings_json, | |
| questions_json, | |
| result_summary_json, | |
| client_json | |
| ) VALUES (?, ?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| run_id, | |
| created_at, | |
| int(config.get("sampleSize") or len(result.get("respondents", [])) or 0), | |
| json.dumps(run_document["respondentSettings"], ensure_ascii=False, sort_keys=True), | |
| json.dumps(run_document["questions"], ensure_ascii=False, sort_keys=True), | |
| json.dumps(run_document["resultSummary"], ensure_ascii=False, sort_keys=True), | |
| json.dumps(client or {}, ensure_ascii=False, sort_keys=True), | |
| ), | |
| ) | |
| info = {"runId": run_id, "createdAt": created_at, "storage": log_status()["storage"], "path": str(path)} | |
| uploader = hf_uploader or upload_run_log_to_hf_dataset | |
| if hf_uploader or hf_log_status()["enabled"]: | |
| try: | |
| info["hfDataset"] = uploader(run_document) | |
| except Exception as exc: # noqa: BLE001 | |
| info["hfDataset"] = {"enabled": True, "uploaded": False, "error": f"{type(exc).__name__}: {exc}"} | |
| else: | |
| info["hfDataset"] = {"enabled": False, "uploaded": False} | |
| return info | |
| def build_run_document(run_id: str, created_at: str, config: dict[str, Any], result: dict[str, Any], client: dict[str, Any]) -> dict[str, Any]: | |
| return { | |
| "runId": run_id, | |
| "createdAt": created_at, | |
| "respondentSettings": build_respondent_settings(config), | |
| "questions": build_question_log(config), | |
| "resultSummary": build_result_summary(result), | |
| "client": client, | |
| } | |
| def recent_runs(limit: int = 20) -> list[dict[str, Any]]: | |
| path = default_log_path() | |
| if not path.exists(): | |
| return [] | |
| ensure_schema(path) | |
| with sqlite3.connect(path) as conn: | |
| conn.row_factory = sqlite3.Row | |
| rows = conn.execute( | |
| """ | |
| SELECT run_id, created_at, sample_size, respondent_settings_json, questions_json, result_summary_json | |
| FROM runs | |
| ORDER BY created_at DESC | |
| LIMIT ? | |
| """, | |
| (max(1, min(100, int(limit))),), | |
| ).fetchall() | |
| return [ | |
| { | |
| "runId": row["run_id"], | |
| "createdAt": row["created_at"], | |
| "sampleSize": row["sample_size"], | |
| "respondentSettings": parse_json(row["respondent_settings_json"]), | |
| "questions": parse_json(row["questions_json"]), | |
| "resultSummary": parse_json(row["result_summary_json"]), | |
| } | |
| for row in rows | |
| ] | |
| def ensure_schema(path: Path) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| with sqlite3.connect(path) as conn: | |
| conn.execute( | |
| """ | |
| CREATE TABLE IF NOT EXISTS runs( | |
| run_id TEXT PRIMARY KEY, | |
| created_at TEXT NOT NULL, | |
| sample_size INTEGER NOT NULL, | |
| respondent_settings_json TEXT NOT NULL, | |
| questions_json TEXT NOT NULL, | |
| result_summary_json TEXT NOT NULL, | |
| client_json TEXT NOT NULL | |
| ) | |
| """ | |
| ) | |
| conn.execute("CREATE INDEX IF NOT EXISTS idx_runs_created_at ON runs(created_at)") | |
| def count_runs(path: Path) -> int: | |
| try: | |
| ensure_schema(path) | |
| with sqlite3.connect(path) as conn: | |
| return int(conn.execute("SELECT COUNT(*) FROM runs").fetchone()[0]) | |
| except sqlite3.Error: | |
| return 0 | |
| def build_respondent_settings(config: dict[str, Any]) -> dict[str, Any]: | |
| return { | |
| "sampleSize": int(config.get("sampleSize") or 0), | |
| "genders": enabled_weighted(config.get("genders")), | |
| "ages": enabled_weighted(config.get("ages")), | |
| "locations": enabled_locations(config), | |
| "personaAttributes": enabled_weighted(config.get("personaAttributes")), | |
| "nemotronFields": [str(item) for item in config.get("nemotronFields", []) if item], | |
| "seed": config.get("seed"), | |
| } | |
| def build_question_log(config: dict[str, Any]) -> list[dict[str, Any]]: | |
| questions = config.get("questions", []) if isinstance(config.get("questions"), list) else [] | |
| logged: list[dict[str, Any]] = [] | |
| for question in questions: | |
| if not isinstance(question, dict): | |
| continue | |
| logged.append( | |
| { | |
| "id": question.get("id"), | |
| "title": question.get("title"), | |
| "source": question.get("source"), | |
| "category": question.get("category"), | |
| "kind": question.get("kind"), | |
| "scale": question.get("scale"), | |
| "lowLabel": question.get("lowLabel"), | |
| "highLabel": question.get("highLabel"), | |
| "options": question.get("options") if question.get("kind") == "categorical" else None, | |
| } | |
| ) | |
| return logged | |
| def build_result_summary(result: dict[str, Any]) -> dict[str, Any]: | |
| return { | |
| "respondents": len(result.get("respondents", [])), | |
| "likertAnswers": len(result.get("likertAnswers", [])), | |
| "categoricalAnswers": len(result.get("categoricalAnswers", [])), | |
| "openAnswers": len(result.get("openAnswers", [])), | |
| "primaryQuestionId": result.get("primaryQuestionId"), | |
| "questionStats": result.get("questionStats", []), | |
| "regionStats": result.get("regionStats", []), | |
| "llmTrace": { | |
| key: value | |
| for key, value in (result.get("llmTrace", {}) if isinstance(result.get("llmTrace"), dict) else {}).items() | |
| if key not in {"rawResponses"} | |
| }, | |
| } | |
| def enabled_weighted(items: Any) -> list[dict[str, Any]]: | |
| return [ | |
| {"id": str(item.get("id")), "weight": item.get("weight")} | |
| for item in items or [] | |
| if isinstance(item, dict) and item.get("enabled") | |
| ] | |
| def enabled_locations(config: dict[str, Any]) -> list[dict[str, Any]]: | |
| labels = { | |
| str(item.get("id")): { | |
| "label": item.get("label"), | |
| "parentRegion": item.get("parentRegion"), | |
| "level": item.get("level"), | |
| } | |
| for item in config.get("locationOptions", []) | |
| if isinstance(item, dict) | |
| } | |
| rows = [] | |
| for item in config.get("locations", []) or []: | |
| if not isinstance(item, dict) or not item.get("enabled"): | |
| continue | |
| location_id = str(item.get("id")) | |
| rows.append({"id": location_id, "weight": item.get("weight"), **labels.get(location_id, {})}) | |
| return rows | |
| def parse_json(value: str) -> Any: | |
| try: | |
| return json.loads(value) | |
| except json.JSONDecodeError: | |
| return None | |