Update hf_api.py

hf_api.py

@@ -17,6 +17,9 @@ from fastapi.responses import JSONResponse
 from response_formatter import ResponseFormatter
 import traceback
 from fastapi.security.api_key import APIKeyHeader, APIKey
+from fastapi.staticfiles import StaticFiles
+from fastapi.openapi.docs import get_swagger_ui_html
+from fastapi.responses import HTMLResponse
 
 # Load environment variables
 load_dotenv()
@@ -300,7 +303,13 @@ class AgentProcessor:
             return None
 
 # Initialize FastAPI app
-app = FastAPI(docs_url="/", redoc_url=None)
+app = FastAPI(
+    title="Agent API",
+    description="API requiring X-API-Key authentication",
+    version="1.0.0",
+    docs_url=None,  # Disable default docs
+    openapi_tags=[{"name": "agent", "description": "Agent endpoints"}],
+)
 agent_processor = None
 
 # Add CORS middleware
@@ -312,6 +321,30 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Add security scheme
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Add security scheme
+app.add_security_requirement({"ApiKeyAuth": []})
+app.openapi_schema = None  # Reset OpenAPI schema
+
+# Define the security scheme
+security_scheme = {
+    "ApiKeyAuth": {
+        "type": "apiKey",
+        "in": "header",
+        "name": "X-API-Key",
+        "description": "API key required for authentication"
+    }
+}
+app.openapi_components = {"securitySchemes": security_scheme}
+
 @app.on_event("startup")
 async def startup_event():
     global agent_processor
@@ -342,10 +375,19 @@ async def get_api_key(
     )
 
 # Update your route to require API key
-@app.post("/v1/agent")
+@app.post(
+    "/v1/agent",
+    tags=["agent"],
+    summary="Process agent request",
+    response_class=StreamingResponse,
+    responses={
+        403: {"description": "Invalid or missing API key"},
+        500: {"description": "Internal server error"}
+    }
+)
 async def process_agent_request(
     request: AgentRequest,
-    api_key: APIKey = Depends(get_api_key)  # Add this line
+    api_key: APIKey = Security(api_key_header, scopes=[])
 ):
     try:
         logger.info(f"Processing agent request: {request.query}")
@@ -375,4 +417,141 @@ if __name__ == "__main__":
         host="0.0.0.0",
         port=port,
         reload=True
+    )
+
+# Add custom docs endpoint
+@app.get("/docs", include_in_schema=False)
+async def custom_swagger_ui_html():
+    return get_swagger_ui_html(
+        openapi_url=app.openapi_url,
+        title=app.title + " - Swagger UI",
+        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
+        swagger_js_url="https://unpkg.com/[email protected]/swagger-ui-bundle.js",
+        swagger_css_url="https://unpkg.com/[email protected]/swagger-ui.css",
+        swagger_favicon_url="https://fastapi.tiangolo.com/img/favicon.png",
+        extra_html="""
+        <style>
+            /* Dark theme with cool colors */
+            :root {
+                --primary-color: #00b4d8;
+                --secondary-color: #90e0ef;
+                --background-color: #0d1117;
+                --text-color: #e6edf3;
+                --border-color: #30363d;
+            }
+            
+            body {
+                background-color: var(--background-color);
+                color: var(--text-color);
+            }
+            
+            .swagger-ui {
+                background-color: var(--background-color);
+                color: var(--text-color);
+            }
+            
+            /* Headers and text */
+            .swagger-ui .info .title,
+            .swagger-ui .info .base-url,
+            .swagger-ui .info li, 
+            .swagger-ui .info p,
+            .swagger-ui .info table {
+                color: var(--text-color);
+            }
+            
+            /* Operation buttons */
+            .swagger-ui .opblock.opblock-post {
+                background: rgba(0, 180, 216, 0.1);
+                border-color: var(--primary-color);
+            }
+            
+            .swagger-ui .opblock.opblock-post .opblock-summary-method {
+                background: var(--primary-color);
+            }
+            
+            /* Authorize button */
+            .swagger-ui .btn.authorize {
+                background: var(--primary-color);
+                border-color: var(--primary-color);
+                color: white;
+            }
+            
+            .swagger-ui .btn.authorize svg {
+                fill: white;
+            }
+            
+            /* Schema sections */
+            .swagger-ui .model-box {
+                background: rgba(48, 54, 61, 0.4);
+            }
+            
+            .swagger-ui .model {
+                color: var(--text-color);
+            }
+            
+            /* Try it out section */
+            .swagger-ui textarea,
+            .swagger-ui input[type=text] {
+                background: var(--background-color);
+                color: var(--text-color);
+                border-color: var(--border-color);
+            }
+            
+            /* Response section */
+            .swagger-ui .responses-table th,
+            .swagger-ui .responses-table td {
+                color: var(--text-color);
+                border-color: var(--border-color);
+            }
+            
+            /* Scrollbar */
+            ::-webkit-scrollbar {
+                width: 8px;
+                height: 8px;
+            }
+            
+            ::-webkit-scrollbar-track {
+                background: var(--background-color);
+            }
+            
+            ::-webkit-scrollbar-thumb {
+                background: var(--primary-color);
+                border-radius: 4px;
+            }
+            
+            /* Code blocks */
+            .swagger-ui .highlight-code {
+                background-color: #1b1f24;
+            }
+            
+            /* Modal dialogs */
+            .swagger-ui .dialog-ux .modal-ux {
+                background: var(--background-color);
+                border-color: var(--border-color);
+            }
+            
+            .swagger-ui .dialog-ux .modal-ux-header h3 {
+                color: var(--text-color);
+            }
+            
+            /* Tables */
+            .swagger-ui table thead tr td,
+            .swagger-ui table thead tr th {
+                color: var(--text-color);
+                border-color: var(--border-color);
+            }
+            
+            /* Links */
+            .swagger-ui a {
+                color: var(--primary-color);
+            }
+            
+            /* Schema dropdowns */
+            .swagger-ui select {
+                background: var(--background-color);
+                color: var(--text-color);
+                border-color: var(--border-color);
+            }
+        </style>
+        """
     )