110 lines
2.8 KiB
Python
110 lines
2.8 KiB
Python
"""FastAPI application for Akern-Genai project.
|
|
|
|
This module provides the web application with WebSocket support
|
|
for streaming responses from the Gemini model.
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
from typing import Annotated
|
|
|
|
from fastapi import (
|
|
FastAPI,
|
|
Request,
|
|
WebSocket,
|
|
WebSocketDisconnect,
|
|
Depends,
|
|
HTTPException,
|
|
status,
|
|
)
|
|
from fastapi.security import HTTPBasic, HTTPBasicCredentials
|
|
from fastapi.templating import Jinja2Templates
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from lib import generate
|
|
|
|
# Configure logging format and level
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Static files configuration
|
|
STATIC_DIR: str = "static"
|
|
TEMPLATES_DIR: str = "templates"
|
|
|
|
# Security configuration
|
|
security = HTTPBasic()
|
|
|
|
|
|
def verify_credentials(credentials: HTTPBasicCredentials = Depends(security)) -> str:
|
|
"""Verify HTTP Basic credentials against environment variables.
|
|
|
|
Args:
|
|
credentials: HTTP Basic authentication credentials.
|
|
|
|
Returns:
|
|
str: The authenticated username.
|
|
|
|
Raises:
|
|
HTTPException: If credentials are invalid.
|
|
"""
|
|
correct_username = os.getenv("BASIC_AUTH_USERNAME")
|
|
correct_password = os.getenv("BASIC_AUTH_PASSWORD")
|
|
|
|
if not (
|
|
credentials.username == correct_username
|
|
and credentials.password == correct_password
|
|
):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Incorrect username or password",
|
|
headers={"WWW-Authenticate": "Basic"},
|
|
)
|
|
return credentials.username
|
|
|
|
|
|
# Initialize FastAPI application
|
|
app = FastAPI()
|
|
app.mount(f"/{STATIC_DIR}", StaticFiles(directory=STATIC_DIR), name="static")
|
|
|
|
templates = Jinja2Templates(directory=os.path.join(STATIC_DIR, TEMPLATES_DIR))
|
|
|
|
|
|
@app.get("/")
|
|
def home(request: Request, username: Annotated[str, Depends(verify_credentials)]):
|
|
"""Render the main index page.
|
|
|
|
Args:
|
|
request: The incoming request object.
|
|
username: The authenticated username from HTTP Basic auth.
|
|
|
|
Returns:
|
|
TemplateResponse: The rendered HTML template.
|
|
"""
|
|
return templates.TemplateResponse("index.html", {"request": request})
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
"""Handle WebSocket connections for streaming responses.
|
|
|
|
Args:
|
|
websocket: The WebSocket connection.
|
|
"""
|
|
await websocket.accept()
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
|
|
async for chunk in generate(data):
|
|
await websocket.send_text(chunk)
|
|
|
|
await websocket.send_text("<<END>>")
|
|
except WebSocketDisconnect:
|
|
logger.info("Client disconnected")
|
|
await websocket.close()
|