Update main.py

main.py

@@ -1,4 +1,3 @@
-# app.py
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any
@@ -7,32 +6,10 @@ import datetime
 
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login, HfHubHTTPError
+from huggingface_hub import login
 
 # ==========================================
-# 1. CONFIGURATION (Secure Defaults)
-# ==========================================
-
-MODEL_ID = "google/functiongemma-270m-it"
-HF_TOKEN_ENV = "HF_TOKEN"
-
-def get_hf_token() -> str:
-    """
-    Fetch Hugging Face token from environment.
-
-    Raises:
-        RuntimeError: if token is missing
-    """
-    token = os.getenv(HF_TOKEN_ENV)
-    if not token:
-        raise RuntimeError(
-            f"Missing required environment variable: {HF_TOKEN_ENV}"
-        )
-    return token
-
-
-# ==========================================
-# 2. APP SETUP
+# 1. APP SETUP
 # ==========================================
 
 app = FastAPI(
@@ -40,45 +17,50 @@ app = FastAPI(
     version="1.0.0",
 )
 
+# Global variables
+MODEL_ID = "google/functiongemma-270m-it"
 tokenizer = None
 model = None
 
-
 # ==========================================
-# 3. DATA MODELS
+# 2. DATA MODELS
 # ==========================================
 
 class ChatRequest(BaseModel):
-    """
-    Request schema for function-call generation.
-    """
     query: str = Field(..., min_length=1, max_length=4096)
     tools: List[Dict[str, Any]]
     include_date: bool = True
 
-
 class HealthResponse(BaseModel):
     status: str
     model: str
-    auth: str
-
+    auth_status: str
 
 # ==========================================
-# 4. STARTUP (Auth + Load Model)
+# 3. STARTUP (Auth + Load Model)
 # ==========================================
 
 @app.on_event("startup")
 async def startup():
     global tokenizer, model
 
-    # A. Authenticate (fail-fast)
+    # A. Authenticate using Environment Variable
+    print("🔐 Checking for HF_TOKEN...")
+    hf_token = os.getenv("HF_TOKEN")
+    
+    if not hf_token:
+        print("❌ Error: HF_TOKEN environment variable is missing.")
+        raise RuntimeError("HF_TOKEN environment variable is missing in Space Settings.")
+    
     try:
-        hf_token = get_hf_token()
         login(token=hf_token)
-    except (RuntimeError, HfHubHTTPError) as e:
-        raise RuntimeError(f"Hugging Face authentication failed: {e}")
+        print("✅ Authentication successful.")
+    except Exception as e:
+        print(f"❌ Authentication Failed: {e}")
+        raise RuntimeError(f"Hugging Face login failed: {e}")
 
     # B. Load Model
+    print(f"🧠 Loading Model: {MODEL_ID}...")
     try:
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
         model = AutoModelForCausalLM.from_pretrained(
@@ -86,12 +68,13 @@ async def startup():
             device_map="cpu",
             torch_dtype=torch.float32,
         )
+        print("✅ Model Loaded Successfully.")
     except Exception as e:
+        print(f"❌ Model Load Failed: {e}")
         raise RuntimeError(f"Model load failed: {e}")
 
-
 # ==========================================
-# 5. API ENDPOINT
+# 4. API ENDPOINT
 # ==========================================
 
 @app.post("/generate")
@@ -137,11 +120,10 @@ async def generate_function_call(request: ChatRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
 @app.get("/", response_model=HealthResponse)
 def health_check():
     return {
         "status": "running",
         "model": MODEL_ID,
-        "auth": "env",
-    }
+        "auth_status": "secure_env",
+    }