Update server.py

server.py

@@ -1,6 +1,8 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from models.transformer.text_generator import TextGenerator
+import os
 
 
 app = FastAPI()
@@ -9,6 +11,18 @@ generator = TextGenerator(
     model_name="frantics-bot-model",
 )
 
+security = HTTPBearer()
+BEARER_TOKEN = os.environ.get("BEARER_TOKEN")
+
+
+async def validate_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    if credentials.credentials == BEARER_TOKEN:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
 
 class Message(BaseModel):
     author: str
@@ -16,16 +30,19 @@ class Message(BaseModel):
 
 
 @app.post("/generate")
-def generate_response(message: Message):
+def generate_response(
+    message: Message,
+    token: str = Depends(validate_token)
+):
     response = generator.generate_text(
         author=message.author,
         input_str=message.content,
         max_length=100,
         num_return_sequences=1,
         do_sample=True,
-        temperature=0.8,  # Слегка уменьшаем уверенность
-        top_k=100,         # Уменьшаем количество рассматриваемых верхних k слов
-        top_p=0.95        # Уменьшаем "ядерность" распределения
+        temperature=0.8,
+        top_k=100,
+        top_p=0.95
     )["generated_texts"][0]
     response = response[:response.find("</s>")]
-    return { "response": response }
+    return {"response": response}