Update app.py

app.py

@@ -14,10 +14,85 @@ import matplotlib.pyplot as plt
 import openai
 from datetime import datetime
 import threading
+from huggingface_hub import hf_hub_download, HfApi
+from huggingface_hub.utils import EntryNotFoundError
+
+USAGE_DATASET_REPO = os.environ.get("USAGE_DATASET_REPO", "NYSERDA-CRE-Working-Group/nyserda_demo_useage_store")
+USAGE_FILENAME = os.environ.get("USAGE_FILENAME", "usage.csv")
+MAX_RUNS_PER_USER = int(os.environ.get("MAX_RUNS_PER_USER", "10"))
 
 os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
 os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+client = OpenAI()
+api = HfApi(token=HF_TOKEN)
+def user_id_from_profile(profile: gr.OAuthProfile | None) -> str | None:
+    if profile is None:
+        return None
+    # You said profile.name exists; normalize it.
+    # If you later can access preferred_username, use that instead (more unique).
+    uid = getattr(profile, "name", None)
+    if not uid:
+        return None
+    return uid.strip().lower()
+
+def _load_usage_df() -> pd.DataFrame:
+    try:
+        local_path = hf_hub_download(
+            repo_id=USAGE_DATASET_REPO,
+            repo_type="dataset",
+            filename=USAGE_FILENAME,
+            token=HF_TOKEN,
+        )
+        return pd.read_csv(local_path)
+    except EntryNotFoundError:
+        # First run: create empty table
+        return pd.DataFrame(columns=["user_id", "runs", "first_seen", "last_seen"])
+
+def _save_usage_df(df: pd.DataFrame, commit_message: str) -> None:
+    tmp_path = "/tmp/usage.csv"
+    df.to_csv(tmp_path, index=False)
+
+    api.upload_file(
+        path_or_fileobj=tmp_path,
+        path_in_repo=USAGE_FILENAME,
+        repo_id=USAGE_DATASET_REPO,
+        repo_type="dataset",
+        commit_message=commit_message,
+    )
+
+def check_and_increment_quota(user_id: str) -> tuple[bool, int]:
+    now = int(time.time())
+    df = _load_usage_df()
+
+    if df.empty or (df["user_id"] == user_id).sum() == 0:
+        runs = 0
+        if runs >= MAX_RUNS_PER_USER:
+            return False, 0
+        new_row = {
+            "user_id": user_id,
+            "runs": 1,
+            "first_seen": now,
+            "last_seen": now,
+        }
+        df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
+        _save_usage_df(df, commit_message=f"usage: increment {user_id} to 1")
+        return True, MAX_RUNS_PER_USER - 1
 
+    idx = df.index[df["user_id"] == user_id][0]
+    runs = int(df.loc[idx, "runs"])
+
+    if runs >= MAX_RUNS_PER_USER:
+        return False, 0
+
+    runs += 1
+    df.loc[idx, "runs"] = runs
+    df.loc[idx, "last_seen"] = now
+
+    _save_usage_df(df, commit_message=f"usage: increment {user_id} to {runs}")
+    return True, MAX_RUNS_PER_USER - runs
+    
 # Global state for the interface
 class InterfaceState:
     def __init__(self):
@@ -1056,7 +1131,20 @@ def agent_worker(user_question):
     state.done = True
 
     
-def run_agentic_workflow(user_question):
+def run_agentic_workflow(user_question, profile: gr.OAuthProfile | None):
+    uid = user_id_from_profile(profile)
+    if uid is None:
+        raise gr.Error("Please sign in with Hugging Face to use this demo.")
+
+    allowed, remaining = check_and_increment_quota(uid)
+    if not allowed:
+        raise gr.Error(f"Usage limit reached: {MAX_RUNS_PER_USER} runs per user.")
+    
+    if remaining <= 2:
+        gr.Warning(f"⚠️ Only {remaining} run(s) left!")
+    else:
+        gr.Info(f"✓ Runs remaining: {remaining}")
+        
     state.done = False
     state.final_answer = ""
 
@@ -1101,6 +1189,8 @@ def run_agentic_workflow(user_question):
 
 # Build Gradio Interface
 with gr.Blocks(title="AEC Compliance Agent") as demo:
+    gr.LoginButton()
+    
     gr.Markdown("# 🏗️ AEC Compliance Analysis Agent")
     gr.Markdown("Ask questions about NYC Building Code compliance for your construction drawings.")
     
@@ -1156,6 +1246,5 @@ with gr.Blocks(title="AEC Compliance Agent") as demo:
 
 if __name__ == "__main__":
     demo.queue().launch(
-        inbrowser=True,
-        auth=("username", os.getenv("PASSWORD"))  # only share the password
+        inbrowser=True
     )