Update main.py

main.py

@@ -1,37 +1,45 @@
-from fastapi import FastAPI, Request, Body
+from fastapi import FastAPI, Request, Body, HTTPException, Depends
+from fastapi.security import APIKeyHeader
+from typing import Optional
 from huggingface_hub import InferenceClient
 import random
 
 API_URL = "https://api-inference.huggingface.co/models/"
+API_KEY = "abcd12345"  # Replace with your actual API key
 
-client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.1"
-)
-
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")
 app = FastAPI()
 
-def format_prompt(message, history):
-  prompt = "<s>"
-  for user_prompt, bot_response in history:
-    prompt += f"[INST] {user_prompt} [/INST]"
-    prompt += f" {bot_response}</s> "
-  prompt += f"[INST] {message} [/INST]"
-  return prompt
+security = APIKeyHeader(name="api_key", auto_error=False)
 
+def get_api_key(api_key: Optional[str] = Depends(security)):
+    if api_key is None or api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Unauthorized access")
+    return api_key
+
+def format_prompt(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
 
 @app.post("/api/v1/generate_text")
-def generate_text(request: Request, prompt: str = Body()):
-    history = []  # You might need to handle this based on your actual usage
-    print(f"request + {request}")
-    temperature = request.headers.get("temperature", 0.5)
-    # print(f"temperature + {temperature}")
-    top_p = request.headers.get("top_p", 0.95)
-    # print(f"top_p + {top_p}")
-    repetition_penalty = request.headers.get("repetition_penalty", 1.0)
-    # print(f"repetition_penalty + {repetition_penalty}")
+def generate_text(
+    request: Request,
+    body: dict = Body(...),
+    api_key: str = Depends(get_api_key)
+):
+    prompt = body.get("prompt", "")
+    sys_prompt = body.get("sysPrompt", "")
+    temperature = body.get("temperature", 0.5)
+    top_p = body.get("top_p", 0.95)
+    repetition_penalty = body.get("repetition_penalty", 1.0)
 
+    history = []  # You might need to handle this based on your actual usage
     formatted_prompt = format_prompt(prompt, history)
-    print(f"formatted_prompt + {formatted_prompt}")
+
     stream = client.text_generation(
         formatted_prompt,
         temperature=temperature,
@@ -41,15 +49,8 @@ def generate_text(request: Request, prompt: str = Body()):
         do_sample=True,
         seed=random.randint(0, 10**7),
         stream=False,
-        details=False, 
+        details=False,
         return_full_text=False
     )
-    # output = ""
-
-    # for response in stream:
-    #     output += response.token.text
-    #     yield output
-        
-    # return output[len(output) - 1]
 
-    return stream
+    return stream