Update handler.py

handler.py

@@ -1,23 +1,21 @@
 # handler.py
 
-import os
 import requests
+import logging
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from typing import Optional
-from dotenv import load_dotenv
 
-# Load secrets from .env
-load_dotenv()
+# Setup logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Securely read environment variables
-HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")
-HUGGINGFACE_MODEL_URL = os.getenv("HUGGINGFACE_MODEL_URL")
-
-# FastAPI router setup
 router = APIRouter()
 
-# Input format
+# Your Hugging Face model URL – must be public
+MODEL_URL = "https://api-inference.huggingface.co/models/CLASSIFIED-HEX/X"
+
+# Input model
 class PromptInput(BaseModel):
     prompt: str
     max_tokens: Optional[int] = 250
@@ -25,17 +23,28 @@ class PromptInput(BaseModel):
     top_p: Optional[float] = 0.95
     top_k: Optional[int] = 50
     repetition_penalty: Optional[float] = 1.2
+    trim_output: Optional[bool] = False  # New feature to remove prompt from result
 
-# Main endpoint
-@router.post("/generate")
-async def generate_text(input_data: PromptInput):
-    if not HUGGINGFACE_API_TOKEN or not HUGGINGFACE_MODEL_URL:
-        raise HTTPException(status_code=500, detail="Hugging Face API token or model URL not configured.")
+# Root health check
+@router.get("/")
+async def root():
+    return {"message": "AI text generation backend is running 🚀"}
 
-    headers = {
-        "Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}"
-    }
+# Ping model check
+@router.get("/ping-model")
+async def ping_model():
+    try:
+        response = requests.post(MODEL_URL, json={"inputs": "ping test"})
+        if response.status_code == 200:
+            return {"status": "Model is online ✅"}
+        else:
+            return {"status": "Model responded with error ❌", "details": response.json()}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Could not reach model: {str(e)}")
 
+# Main generation route
+@router.post("/generate")
+async def generate_text(input_data: PromptInput):
     payload = {
         "inputs": input_data.prompt,
         "parameters": {
@@ -48,18 +57,25 @@ async def generate_text(input_data: PromptInput):
     }
 
     try:
-        response = requests.post(HUGGINGFACE_MODEL_URL, headers=headers, json=payload)
+        logger.info(f"Sending prompt to model: {input_data.prompt}")
+        response = requests.post(MODEL_URL, json=payload)
 
         if response.status_code != 200:
+            logger.error(f"Model error: {response.status_code} - {response.text}")
             raise HTTPException(status_code=response.status_code, detail=response.json())
 
         result = response.json()
-        generated_text = result[0].get("generated_text") if isinstance(result, list) else result.get("generated_text", "")
+        raw_output = result[0].get("generated_text") if isinstance(result, list) else result.get("generated_text", "")
+
+        # Optionally trim prompt from beginning
+        if input_data.trim_output and raw_output.startswith(input_data.prompt):
+            raw_output = raw_output[len(input_data.prompt):].lstrip()
 
         return {
             "status": "success",
-            "output": generated_text
+            "output": raw_output
         }
 
     except Exception as e:
+        logger.exception("Text generation failed")
         raise HTTPException(status_code=500, detail=f"Text generation failed: {str(e)}")