updates app

app.py

@@ -1,11 +1,14 @@
 import asyncio
 import time
-from fastapi import FastAPI, Request, HTTPException
+import nh3
+from auth import APIKeyMiddleware
+from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 from starlette.status import HTTP_504_GATEWAY_TIMEOUT
 from neural_searcher import NeuralSearcher
 from huggingface_hub import login
 from config import HUGGING_FACE_API_KEY, COLLECTION_NAME, API_KEY
+from fastapi.security.api_key import APIKey
 
 login(HUGGING_FACE_API_KEY)
 
@@ -14,31 +17,28 @@ app = FastAPI()
 neural_searcher = NeuralSearcher(collection_name=COLLECTION_NAME)
 
 REQUEST_TIMEOUT_ERROR = 30
-ALLOWED_API_KEY = API_KEY
 
+app.add_middleware(APIKeyMiddleware)
 
 @app.get("/api/search")
 async def search(q: str):
-    data = await neural_searcher.search(text=q)
+    xss = nh3.clean(q)
+    data = await neural_searcher.search(text=xss)
     return data
 
 
-@app.middleware("http")
-async def combined_middleware(request: Request, call_next):
-    start_time = time.time()
-
-    api_key = request.headers.get("X-API-KEY", "")
-    if api_key != ALLOWED_API_KEY:
-        return JSONResponse(status_code=403, content={"detail": "Access denied."})
-
-    try:
-        return await asyncio.wait_for(call_next(request), timeout=REQUEST_TIMEOUT_ERROR)
-    except asyncio.TimeoutError:
-        process_time = time.time() - start_time
-        return JSONResponse(
-            content={
-                "detail": "Request processing time exceeded limit",
-                "processing_time": process_time
-            },
-            status_code=HTTP_504_GATEWAY_TIMEOUT
-        )
+# @app.middleware("http")
+# async def combined_middleware(request: Request, call_next):
+#     start_time = time.time()
+
+#     try:
+#         return await asyncio.wait_for(call_next(request), timeout=REQUEST_TIMEOUT_ERROR)
+#     except asyncio.TimeoutError:
+#         process_time = time.time() - start_time
+#         return JSONResponse(
+#             content={
+#                 "detail": "Request processing time exceeded limit",
+#                 "processing_time": process_time
+#             },
+#             status_code=HTTP_504_GATEWAY_TIMEOUT
+#         )

auth.py

@@ -0,0 +1,12 @@
+from config import API_KEY
+
+from fastapi import Request, HTTPException
+from starlette.middleware.base import BaseHTTPMiddleware
+
+class APIKeyMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        api_key = request.headers.get("X-API-KEY")
+        if api_key != API_KEY:
+            raise HTTPException(status_code=403, detail="Forbidden")
+        response = await call_next(request)
+        return response

requirements.txt

@@ -8,3 +8,5 @@ python-dotenv
 qdrant-client
 qdrant-client[fastembed]>=1.8.2
 sentence-transformers
+fastapi-csrf-protect
+nh3