api key added

app.py

@@ -1,6 +1,7 @@
 import os
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
 from pydantic import BaseModel
+from fastapi.security.api_key import APIKeyHeader
 from utils import retrive_context, generate_response
 
 
@@ -11,32 +12,35 @@ app = FastAPI()
 class QueryRequest(BaseModel):
     # Asked query should be in string format
     query: str
-    api_key: str             
-
+    
 
 class QueryResponse(BaseModel):
     # Response should be in string format
     response: str
 
 
+api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
+
+def get_api_key(api_key: str = Depends(api_key_header)):
+    if api_key == os.getenv("API_KEY"):
+        return api_key
+    else:
+        raise HTTPException(status_code=401, detail="Invalid API key")
 
 @app.post("/infer", response_model=QueryResponse)
-def infer(query_request: QueryRequest):
-    key = query_request.api_key
+def infer(query_request: QueryRequest, api_key: str = Depends(get_api_key)):
     query = query_request.query
-    
-    if (key == os.getenv("API_KEY")):      
-        context = retrive_context(query)
-        if context == 500:
-            raise HTTPException(status_code=500, detail="Error retrieving context")
-    
-        response = generate_response(query, context)
-        if response == 500:
-            raise HTTPException(status_code=500, detail="Error generating response")
-    
-        return QueryResponse(response=response)
-    else:
-        raise HTTPException(status_code=401, detail="Invalid api key")
+
+    context = retrive_context(query)
+    if context == 500:
+        raise HTTPException(status_code=500, detail="Error retrieving context")
+
+    response = generate_response(query, context)
+    if response == 500:
+        raise HTTPException(status_code=500, detail="Error generating response")
+
+    return QueryResponse(response=response)
+
 
 # Root endpoint for testing
 @app.get("/")