Update hf_api.py

hf_api.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Security, Depends
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Dict, List, Optional, Union, Any
 from pydantic import BaseModel, Field
@@ -16,6 +16,7 @@ from fastapi.responses import StreamingResponse
 from fastapi.responses import JSONResponse
 from response_formatter import ResponseFormatter
 import traceback
+from fastapi.security.api_key import APIKeyHeader, APIKey
 
 # Load environment variables
 load_dotenv()
@@ -27,6 +28,11 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Add these constants near the top of the file after imports
+API_KEY_NAME = "X-API-Key"
+API_KEY = os.getenv("CLIENT_API_KEY")  # Add this to your .env file
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=True)
+
 class AgentOutput(BaseModel):
     """Structured output from agent processing"""
     thought_content: str
@@ -294,7 +300,7 @@ class AgentProcessor:
             return None
 
 # Initialize FastAPI app
-app = FastAPI()
+app = FastAPI(docs_url="/", redoc_url=None)
 agent_processor = None
 
 # Add CORS middleware
@@ -318,12 +324,32 @@ async def shutdown_event():
     if agent_processor:
         await agent_processor.cleanup()
 
+# Add this function before your routes
+async def get_api_key(
+    api_key_header: str = Security(api_key_header)
+) -> APIKey:
+    """Validate API key from header."""
+    if not API_KEY:
+        raise HTTPException(
+            status_code=500,
+            detail="API key configuration is missing on server"
+        )
+    if api_key_header == API_KEY:
+        return api_key_header
+    raise HTTPException(
+        status_code=403,
+        detail="Invalid or missing API key"
+    )
+
+# Update your route to require API key
 @app.post("/v1/agent")
-async def process_agent_request(request: AgentRequest):
+async def process_agent_request(
+    request: AgentRequest,
+    api_key: APIKey = Depends(get_api_key)  # Add this line
+):
     try:
         logger.info(f"Processing agent request: {request.query}")
         return await agent_processor.process_stream(request)
-        
     except Exception as e:
         logger.error(f"Error in agent request processing: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=str(e))