Spaces:
Sleeping
Sleeping
Update hf_api.py
Browse files
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(
|
|
|
|
|
|
|
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))
|