Update main.py

main.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
@@ -9,15 +10,16 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
-# Initialize FastAPI app
+# Initialize FastAPI app and security
 app = FastAPI()
+security = HTTPBearer()
 
-# Get Hugging Face token from environment variable
+# Get HuggingFace token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
-    raise ValueError("HF_TOKEN environment variable not set")
+    raise ValueError("HF_TOKEN environment variable is not set")
 
-# Initialize Hugging Face client with token
+# Initialize HuggingFace client with token
 client = InferenceClient(
     model="mistralai/Mixtral-8x7B-Instruct-v0.1",
     token=HF_TOKEN
@@ -36,10 +38,10 @@ class GenerationRequest(BaseModel):
     top_p: Optional[float] = 0.95
 
 def format_prompt(message: str, history: List[ChatMessage] = None, system_message: str = None) -> str:
-    prompt = "<s>"
+    prompt = ""
     
     if system_message:
-        prompt += f"[INST] {system_message} [/INST]</s>"
+        prompt += f"<s>[INST] {system_message} [/INST]</s>"
     
     if history:
         for msg in history:
@@ -51,37 +53,59 @@ def format_prompt(message: str, history: List[ChatMessage] = None, system_messag
     prompt += f"<s>[INST] {message} [/INST]"
     return prompt
 
+async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    if credentials.credentials != HF_TOKEN:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid authentication credentials"
+        )
+    return credentials.credentials
+
 @app.post("/generate/")
-async def generate_text(request: GenerationRequest):
+async def generate_text(
+    request: GenerationRequest,
+    token: str = Depends(verify_token)
+):
     try:
         message = request.prompt if request.prompt else request.message
         if not message:
-            raise HTTPException(status_code=400, detail="Either 'prompt' or 'message' must be provided")
+            raise HTTPException(
+                status_code=400, 
+                detail="Either 'prompt' or 'message' must be provided"
+            )
 
-        # Format the prompt
         formatted_prompt = format_prompt(
             message=message,
             history=request.history,
             system_message=request.system_message
         )
 
-        # Make the request to Hugging Face
+        parameters = {
+            "temperature": max(request.temperature, 0.01),
+            "top_p": request.top_p,
+            "max_new_tokens": 1048,
+            "do_sample": True,
+            "return_full_text": False
+        }
+
         response = client.text_generation(
             formatted_prompt,
-            max_new_tokens=1024,
-            temperature=max(request.temperature, 0.01),
-            top_p=request.top_p,
-            do_sample=True,
-            seed=42
+            **parameters
         )
-        
+
         if not response:
-            raise HTTPException(status_code=500, detail="Empty response from model")
+            raise HTTPException(
+                status_code=500,
+                detail="No response received from model"
+            )
 
         return {"response": response}
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error generating response: {str(e)}"
+        )
 
 @app.get("/health")
 async def health_check():