api key added

app.py

@@ -1,3 +1,4 @@
+import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from utils import retrive_context, generate_response
@@ -9,7 +10,8 @@ app = FastAPI()
 
 class QueryRequest(BaseModel):
     # Asked query should be in string format
-    query: str             
+    query: str
+    api_key: str             
 
 
 class QueryResponse(BaseModel):
@@ -20,17 +22,21 @@ class QueryResponse(BaseModel):
 
 @app.post("/infer", response_model=QueryResponse)
 def infer(query_request: QueryRequest):
+    key = query_request.api_key
     query = query_request.query
-    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")
+    if (key == os.getenv("API_KEY")):      
+        context = retrive_context(query)
+        if context == 500:
+            raise HTTPException(status_code=500, detail="Error retrieving context")
     
-    return QueryResponse(response=response)
-
+        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")
 
 # Root endpoint for testing
 @app.get("/")