Severian commited on
Commit
73b0144
·
verified ·
1 Parent(s): 3dc43a2

Update hf_api.py

Browse files
Files changed (1) hide show
  1. hf_api.py +30 -4
hf_api.py CHANGED
@@ -1,4 +1,4 @@
1
- from fastapi import FastAPI, HTTPException, Request
2
  from fastapi.middleware.cors import CORSMiddleware
3
  from typing import Dict, List, Optional, Union, Any
4
  from pydantic import BaseModel, Field
@@ -16,6 +16,7 @@ from fastapi.responses import StreamingResponse
16
  from fastapi.responses import JSONResponse
17
  from response_formatter import ResponseFormatter
18
  import traceback
 
19
 
20
  # Load environment variables
21
  load_dotenv()
@@ -27,6 +28,11 @@ logging.basicConfig(
27
  )
28
  logger = logging.getLogger(__name__)
29
 
 
 
 
 
 
30
  class AgentOutput(BaseModel):
31
  """Structured output from agent processing"""
32
  thought_content: str
@@ -294,7 +300,7 @@ class AgentProcessor:
294
  return None
295
 
296
  # Initialize FastAPI app
297
- app = FastAPI()
298
  agent_processor = None
299
 
300
  # Add CORS middleware
@@ -318,12 +324,32 @@ async def shutdown_event():
318
  if agent_processor:
319
  await agent_processor.cleanup()
320
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
321
  @app.post("/v1/agent")
322
- async def process_agent_request(request: AgentRequest):
 
 
 
323
  try:
324
  logger.info(f"Processing agent request: {request.query}")
325
  return await agent_processor.process_stream(request)
326
-
327
  except Exception as e:
328
  logger.error(f"Error in agent request processing: {e}", exc_info=True)
329
  raise HTTPException(status_code=500, detail=str(e))
 
1
+ from fastapi import FastAPI, HTTPException, Request, Security, Depends
2
  from fastapi.middleware.cors import CORSMiddleware
3
  from typing import Dict, List, Optional, Union, Any
4
  from pydantic import BaseModel, Field
 
16
  from fastapi.responses import JSONResponse
17
  from response_formatter import ResponseFormatter
18
  import traceback
19
+ from fastapi.security.api_key import APIKeyHeader, APIKey
20
 
21
  # Load environment variables
22
  load_dotenv()
 
28
  )
29
  logger = logging.getLogger(__name__)
30
 
31
+ # Add these constants near the top of the file after imports
32
+ API_KEY_NAME = "X-API-Key"
33
+ API_KEY = os.getenv("CLIENT_API_KEY") # Add this to your .env file
34
+ api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=True)
35
+
36
  class AgentOutput(BaseModel):
37
  """Structured output from agent processing"""
38
  thought_content: str
 
300
  return None
301
 
302
  # Initialize FastAPI app
303
+ app = FastAPI(docs_url="/", redoc_url=None)
304
  agent_processor = None
305
 
306
  # Add CORS middleware
 
324
  if agent_processor:
325
  await agent_processor.cleanup()
326
 
327
+ # Add this function before your routes
328
+ async def get_api_key(
329
+ api_key_header: str = Security(api_key_header)
330
+ ) -> APIKey:
331
+ """Validate API key from header."""
332
+ if not API_KEY:
333
+ raise HTTPException(
334
+ status_code=500,
335
+ detail="API key configuration is missing on server"
336
+ )
337
+ if api_key_header == API_KEY:
338
+ return api_key_header
339
+ raise HTTPException(
340
+ status_code=403,
341
+ detail="Invalid or missing API key"
342
+ )
343
+
344
+ # Update your route to require API key
345
  @app.post("/v1/agent")
346
+ async def process_agent_request(
347
+ request: AgentRequest,
348
+ api_key: APIKey = Depends(get_api_key) # Add this line
349
+ ):
350
  try:
351
  logger.info(f"Processing agent request: {request.query}")
352
  return await agent_processor.process_stream(request)
 
353
  except Exception as e:
354
  logger.error(f"Error in agent request processing: {e}", exc_info=True)
355
  raise HTTPException(status_code=500, detail=str(e))