import os import torch from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Depends, BackgroundTasks, Request, Query, APIRouter, Path, Body, status, Response, Header from fastapi.responses import StreamingResponse, JSONResponse, FileResponse, HTMLResponse, PlainTextResponse, RedirectResponse from pydantic import BaseModel, validator, Field, root_validator, EmailStr, constr, ValidationError from transformers import ( AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList, pipeline, AutoProcessor, AutoModelForImageClassification, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering, AutoModelForSpeechSeq2Seq, AutoModelForImageSegmentation, AutoFeatureExtractor, AutoModelForTokenClassification, AutoModelForMaskedLM, AutoModelForImageInpainting, AutoModelForImageSuperResolution, AutoModelForObjectDetection, AutoModelForImageCaptioning, AutoModelForTextToSpeech, AutoModelForSeq2SeqLM ) from io import BytesIO import boto3 from botocore.exceptions import NoCredentialsError, ClientError from huggingface_hub import snapshot_download import asyncio import tempfile import hashlib from PIL import Image import base64 from typing import Optional, List, Union, Dict, Any import uuid import subprocess import json from starlette.middleware.cors import CORSMiddleware import numpy as np from typing import Dict, Any from fastapi.staticfiles import StaticFiles from fastapi.templating import Jinja2Templates from fastapi.middleware.gzip import GZipMiddleware from transformers import AutoImageProcessor, pipeline from fastapi.security import APIKeyHeader, OAuth2PasswordBearer, OAuth2PasswordRequestForm from fastapi.security.api_key import APIKeyCookie from fastapi import Depends, Security, status, APIRouter, UploadFile, File, Request from fastapi.security import APIKeyHeader, OAuth2PasswordRequestForm from passlib.context import CryptContext from jose import JWTError, jwt from datetime import datetime, timedelta from starlette.requests import Request import logging from pydantic import EmailStr, constr, ValidationError from database import insert_user, get_user, delete_user, update_user, create_db_and_table from starlette.middleware import Middleware from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.types import ASGIApp import uvicorn from starlette.responses import StreamingResponse import logging from pydantic import EmailStr, constr, ValidationError from database import insert_user, get_user, delete_user, update_user, create_db_and_table, get_all_users from starlette.middleware import Middleware from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.types import ASGIApp import uvicorn from starlette.responses import StreamingResponse import logging from fastapi.exceptions import RequestValidationError from fastapi import Request, status, Depends from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer from jose import JWTError, jwt from passlib.context import CryptContext from datetime import datetime, timedelta from typing import Optional #setting up logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(filename)s - %(lineno)d - %(message)s') logger = logging.getLogger(__name__) #JWT Settings SECRET_KEY = os.getenv("SECRET_KEY") if not SECRET_KEY: raise ValueError("SECRET_KEY must be set.") ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 30 #Password Hashing pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") #Database connection - replace with your database setup #Example using SQLite import sqlite3 conn = sqlite3.connect('users.db') cursor = conn.cursor() #OAuth2 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") #API Key API_KEY = os.getenv("API_KEY") api_key_header = APIKeyHeader(name="X-API-Key") #Configuration AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID") AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY") AWS_REGION = os.getenv("AWS_REGION") S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME") HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") TEMP_DIR = "/tmp" STATIC_DIR = "static" TEMPLATES = Jinja2Templates(directory="templates") app = FastAPI() app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") app.add_middleware(GZipMiddleware) origins = ["*"] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class User(BaseModel): username: constr(min_length=3, max_length=50) email: EmailStr password: constr(min_length=8) class GenerateRequest(BaseModel): model_name: str input_text: Optional[str] = Field(None, description="Input text for generation.") task_type: str = Field(..., description="Type of generation task (text, image, audio, video, classification, translation, question-answering, speech-to-text, text-to-speech, image-segmentation, feature-extraction, token-classification, fill-mask, image-inpainting, image-super-resolution, object-detection, image-captioning, audio-transcription, summarization).") temperature: float = 1.0 max_new_tokens: int = 200 stream: bool = True top_p: float = 1.0 top_k: int = 50 repetition_penalty: float = 1.0 num_return_sequences: int = 1 do_sample: bool = True chunk_delay: float = 0.0 stop_sequences: List[str] = [] image_file: Optional[UploadFile] = None source_language: Optional[str] = None target_language: Optional[str] = None context: Optional[str] = None audio_file: Optional[UploadFile] = None raw_input: Optional[Union[str, bytes]] = None # for feature extraction masked_text: Optional[str] = None # for fill-mask mask_image: Optional[UploadFile] = None # for image inpainting low_res_image: Optional[UploadFile] = None # for image super-resolution @validator("task_type") def validate_task_type(cls, value): allowed_types = ["text", "image", "audio", "video", "classification", "translation", "question-answering", "speech-to-text", "text-to-speech", "image-segmentation", "feature-extraction", "token-classification", "fill-mask", "image-inpainting", "image-super-resolution", "object-detection", "image-captioning", "audio-transcription", "summarization"] if value not in allowed_types: raise ValueError(f"Invalid task_type. Allowed types are: {allowed_types}") return value @root_validator def check_input(cls, values): task_type = values.get("task_type") if task_type == "text" and values.get("input_text") is None: raise ValueError("input_text is required for text generation.") elif task_type == "speech-to-text" and values.get("audio_file") is None: raise ValueError("audio_file is required for speech-to-text.") elif task_type == "classification" and values.get("image_file") is None: raise ValueError("image_file is required for image classification.") elif task_type == "image-segmentation" and values.get("image_file") is None: raise ValueError("image_file is required for image segmentation.") elif task_type == "feature-extraction" and values.get("raw_input") is None: raise ValueError("raw_input is required for feature extraction.") elif task_type == "fill-mask" and values.get("masked_text") is None: raise ValueError("masked_text is required for fill-mask.") elif task_type == "image-inpainting" and (values.get("image_file") is None or values.get("mask_image") is None): raise ValueError("image_file and mask_image are required for image inpainting.") elif task_type == "image-super-resolution" and values.get("low_res_image") is None: raise ValueError("low_res_image is required for image super-resolution.") return values class S3ModelLoader: def __init__(self, bucket_name, aws_access_key_id, aws_secret_access_key, aws_region): self.bucket_name = bucket_name self.s3 = boto3.client( 's3', aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, region_name=aws_region ) def _get_s3_uri(self, model_name): return f"{self.bucket_name}/{model_name.replace('/', '-')}" def load_model_and_tokenizer(self, model_name, task_type): s3_uri = self._get_s3_uri(model_name) try: self.s3.head_object(Bucket=self.bucket_name, Key=f'{s3_uri}/config.json') except ClientError as e: if e.response['Error']['Code'] == '404': with tempfile.TemporaryDirectory() as tmpdir: model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN, cache_dir=tmpdir) self._upload_model_to_s3(model_path, s3_uri) else: raise HTTPException(status_code=500, detail=f"Error accessing S3: {e}") return self._load_from_s3(s3_uri, task_type) def _upload_model_to_s3(self, model_path, s3_uri): for root, _, files in os.walk(model_path): for file in files: local_path = os.path.join(root, file) s3_path = os.path.join(s3_uri, os.path.relpath(local_path, model_path)) self.s3.upload_file(local_path, self.bucket_name, s3_path) def _load_from_s3(self, s3_uri, task_type): with tempfile.TemporaryDirectory() as tmpdir: model_path = os.path.join(tmpdir, s3_uri) os.makedirs(model_path, exist_ok=True) self.s3.download_file(self.bucket_name, f"{s3_uri}/config.json", os.path.join(model_path, "config.json")) if task_type == "text": model = AutoModelForCausalLM.from_pretrained(model_path, load_in_8bit=True) tokenizer = AutoTokenizer.from_pretrained(model_path) if tokenizer.eos_token_id is None: tokenizer.eos_token_id = tokenizer.pad_token_id return {"model": model, "tokenizer": tokenizer, "pad_token_id": tokenizer.pad_token_id, "eos_token_id": tokenizer.eos_token_id} elif task_type in ["image", "audio", "video"]: processor = AutoProcessor.from_pretrained(model_path) pipeline_function = pipeline(task_type, model=model_path, device=0 if torch.cuda.is_available() else -1, processor=processor) return {"pipeline": pipeline_function} elif task_type == "classification": model = AutoModelForImageClassification.from_pretrained(model_path) processor = AutoProcessor.from_pretrained(model_path) return {"model": model, "processor": processor} elif task_type == "translation": model = AutoModelForSeq2SeqLM.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) return {"model": model, "tokenizer": tokenizer} elif task_type == "question-answering": model = AutoModelForQuestionAnswering.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) return {"model": model, "tokenizer": tokenizer} elif task_type == "speech-to-text": model = pipeline("automatic-speech-recognition", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "text-to-speech": model = pipeline("text-to-speech", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "image-segmentation": model = pipeline("image-segmentation", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "feature-extraction": feature_extractor = AutoFeatureExtractor.from_pretrained(model_path) return {"feature_extractor": feature_extractor} elif task_type == "token-classification": model = AutoModelForTokenClassification.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) return {"model": model, "tokenizer": tokenizer} elif task_type == "fill-mask": model = AutoModelForMaskedLM.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) return {"model": model, "tokenizer": tokenizer} elif task_type == "image-inpainting": model = pipeline("image-inpainting", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "image-super-resolution": model = pipeline("image-super-resolution", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "object-detection": model = pipeline("object-detection", model=model_path, device=0 if torch.cuda.is_available() else -1) image_processor = AutoImageProcessor.from_pretrained(model_path) return {"pipeline": model, "image_processor": image_processor} elif task_type == "image-captioning": model = pipeline("image-captioning", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "audio-transcription": model = pipeline("automatic-speech-recognition", model=model_path, device=0 if torch.cuda.is_available() else -1) return {"pipeline": model} elif task_type == "summarization": model = pipeline("summarization", model=model_path, device=0 if torch.cuda.is_available() else -1) tokenizer = AutoTokenizer.from_pretrained(model_path) return {"model": model, "tokenizer": tokenizer} else: raise ValueError("Unsupported task type") async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay): encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True).to(device) input_length = encoded_input["input_ids"].shape[1] max_length = model.config.max_length remaining_tokens = max_length - input_length if remaining_tokens <= 0: yield "" generation_config.max_new_tokens = min(remaining_tokens, generation_config.max_new_tokens) def stop_criteria(input_ids, scores): decoded_output = tokenizer.decode(input_ids[0][-1], skip_special_tokens=True) return decoded_output in stop_sequences stopping_criteria = StoppingCriteriaList([stop_criteria]) outputs = model.generate( **encoded_input, do_sample=generation_config.do_sample, max_new_tokens=generation_config.max_new_tokens, temperature=generation_config.temperature, top_p=generation_config.top_p, top_k=generation_config.top_k, repetition_penalty=generation_config.repetition_penalty, num_return_sequences=generation_config.num_return_sequences, stopping_criteria=stopping_criteria, output_scores=True, return_dict_in_generate=True ) for output in outputs.sequences: for token_id in output: token = tokenizer.decode(token_id, skip_special_tokens=True) yield token model_loader = S3ModelLoader(S3_BUCKET_NAME, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION) def get_model_data(request: GenerateRequest): return model_loader.load_model_and_tokenizer(request.model_name, request.task_type) async def verify_api_key(api_key: str = Depends(api_key_header)): if api_key != API_KEY: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key") @app.post("/generate", dependencies=[Depends(verify_api_key)]) async def generate(request: GenerateRequest, background_tasks: BackgroundTasks, model_data = Depends(get_model_data)): try: device = "cuda" if torch.cuda.is_available() else "cpu" if request.task_type == "text": model = model_data["model"].to(device) tokenizer = model_data["tokenizer"] generation_config = GenerationConfig( temperature=request.temperature, max_new_tokens=request.max_new_tokens, top_p=request.top_p, top_k=request.top_k, repetition_penalty=request.repetition_penalty, do_sample=request.do_sample, num_return_sequences=request.num_return_sequences, ) async def stream_with_tokens(): async for token in stream_text(model, tokenizer, request.input_text, generation_config, request.stop_sequences, device, request.chunk_delay): yield f"Token: {token}\n" return StreamingResponse(stream_with_tokens(), media_type="text/plain") elif request.task_type in ["image", "audio", "video"]: pipeline = model_data["pipeline"] result = pipeline(request.input_text) if request.task_type == "image": image = result[0] img_byte_arr = BytesIO() image.save(img_byte_arr, format="PNG") img_byte_arr.seek(0) return StreamingResponse(img_byte_arr, media_type="image/png") elif request.task_type == "audio": audio = result[0] audio_byte_arr = BytesIO() audio.save(audio_byte_arr, format="wav") audio_byte_arr.seek(0) return StreamingResponse(audio_byte_arr, media_type="audio/wav") elif request.task_type == "video": video = result[0] video_byte_arr = BytesIO() video.save(video_byte_arr, format="mp4") video_byte_arr.seek(0) return StreamingResponse(video_byte_arr, media_type="video/mp4") elif request.task_type == "classification": if request.image_file is None: raise HTTPException(status_code=400, detail="Image file is required for classification.") contents = await request.image_file.read() image = Image.open(BytesIO(contents)).convert("RGB") model = model_data["model"].to(device) processor = model_data["processor"] inputs = processor(images=image, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) predicted_class_idx = outputs.logits.argmax().item() predicted_class = model.config.id2label[predicted_class_idx] return JSONResponse({"predicted_class": predicted_class}) elif request.task_type == "translation": if request.source_language is None or request.target_language is None: raise HTTPException(status_code=400, detail="Source and target languages are required for translation.") model = model_data["model"].to(device) tokenizer = model_data["tokenizer"] inputs = tokenizer(request.input_text, return_tensors="pt").to(device) with torch.no_grad(): outputs = model.generate(**inputs) translation = tokenizer.decode(outputs[0], skip_special_tokens=True) return JSONResponse({"translation": translation}) elif request.task_type == "question-answering": if request.context is None: raise HTTPException(status_code=400, detail="Context is required for question answering.") model = model_data["model"].to(device) tokenizer = model_data["tokenizer"] inputs = tokenizer(question=request.input_text, context=request.context, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) answer_start = torch.argmax(outputs.start_logits) answer_end = torch.argmax(outputs.end_logits) + 1 answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end])) return JSONResponse({"answer": answer}) elif request.task_type == "speech-to-text": if request.audio_file is None: raise HTTPException(status_code=400, detail="Audio file is required for speech-to-text.") contents = await request.audio_file.read() pipeline = model_data["pipeline"] try: transcription = pipeline(contents, sampling_rate=16000)[0]["text"] # Assuming 16kHz sampling rate return JSONResponse({"transcription": transcription}) except Exception as e: raise HTTPException(status_code=500, detail=f"Error during speech-to-text: {str(e)}") elif request.task_type == "text-to-speech": if not request.input_text: raise HTTPException(status_code=400, detail="Input text is required for text-to-speech.") pipeline = model_data["pipeline"] try: audio = pipeline(request.input_text)[0] file_path = os.path.join(TEMP_DIR, f"{uuid.uuid4()}.wav") audio.save(file_path) background_tasks.add_task(os.remove, file_path) return FileResponse(file_path, media_type="audio/wav") except Exception as e: raise HTTPException(status_code=500, detail=f"Error during text-to-speech: {str(e)}") elif request.task_type == "image-segmentation": if request.image_file is None: raise HTTPException(status_code=400, detail="Image file is required for image segmentation.") contents = await request.image_file.read() image = Image.open(BytesIO(contents)).convert("RGB") pipeline = model_data["pipeline"] result = pipeline(image) mask = result[0]['mask'] mask_byte_arr = BytesIO() mask.save(mask_byte_arr, format="PNG") mask_byte_arr.seek(0) return StreamingResponse(mask_byte_arr, media_type="image/png") elif request.task_type == "feature-extraction": if request.raw_input is None: raise HTTPException(status_code=400, detail="raw_input is required for feature extraction.") feature_extractor = model_data["feature_extractor"] try: if isinstance(request.raw_input, str): inputs = feature_extractor(text=request.raw_input, return_tensors="pt") elif isinstance(request.raw_input, bytes): image = Image.open(BytesIO(request.raw_input)).convert("RGB") inputs = feature_extractor(images=image, return_tensors="pt") else: raise ValueError("Unsupported raw_input type.") features = inputs.pixel_values # Adjust according to your feature extractor return JSONResponse({"features": features.tolist()}) except Exception as fe: raise HTTPException(status_code=400, detail=f"Error during feature extraction: {fe}") elif request.task_type == "token-classification": if request.input_text is None: raise HTTPException(status_code=400, detail="Input text is required for token classification.") model = model_data["model"].to(device) tokenizer = model_data["tokenizer"] inputs = tokenizer(request.input_text, return_tensors="pt", padding=True, truncation=True) with torch.no_grad(): outputs = model(**inputs) predictions = outputs.logits.argmax(dim=-1) predicted_labels = [model.config.id2label[label_id] for label_id in predictions[0].tolist()] return JSONResponse({"predicted_labels": predicted_labels}) elif request.task_type == "fill-mask": if request.masked_text is None: raise HTTPException(status_code=400, detail="masked_text is required for fill-mask.") model = model_data["model"].to(device) tokenizer = model_data["tokenizer"] inputs = tokenizer(request.masked_text, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits masked_index = torch.where(inputs.input_ids == tokenizer.mask_token_id)[1] predicted_token_id = torch.argmax(logits[0, masked_index]) predicted_token = tokenizer.decode(predicted_token_id) return JSONResponse({"predicted_token": predicted_token}) elif request.task_type == "image-inpainting": if request.image_file is None or request.mask_image is None: raise HTTPException(status_code=400, detail="image_file and mask_image are required for image inpainting.") image_contents = await request.image_file.read() mask_contents = await request.mask_image.read() image = Image.open(BytesIO(image_contents)).convert("RGB") mask = Image.open(BytesIO(mask_contents)).convert("L") # Assuming mask is grayscale pipeline = model_data["pipeline"] result = pipeline(image, mask) inpainted_image = result[0] img_byte_arr = BytesIO() inpainted_image.save(img_byte_arr, format="PNG") img_byte_arr.seek(0) return StreamingResponse(img_byte_arr, media_type="image/png") elif request.task_type == "image-super-resolution": if request.low_res_image is None: raise HTTPException(status_code=400, detail="low_res_image is required for image super-resolution.") contents = await request.low_res_image.read() image = Image.open(BytesIO(contents)).convert("RGB") pipeline = model_data["pipeline"] result = pipeline(image) upscaled_image = result[0] img_byte_arr = BytesIO() upscaled_image.save(img_byte_arr, format="PNG") img_byte_arr.seek(0) return StreamingResponse(img_byte_arr, media_type="image/png") elif request.task_type == "object-detection": if request.image_file is None: raise HTTPException(status_code=400, detail="Image file is required for object detection.") contents = await request.image_file.read() image = Image.open(BytesIO(contents)).convert("RGB") pipeline = model_data["pipeline"] image_processor = model_data["image_processor"] inputs = image_processor(images=image, return_tensors="pt") with torch.no_grad(): outputs = pipeline(image) detections = outputs return JSONResponse({"detections": detections}) elif request.task_type == "image-captioning": if request.image_file is None: raise HTTPException(status_code=400, detail="Image file is required for image captioning.") contents = await request.image_file.read() image = Image.open(BytesIO(contents)).convert("RGB") pipeline = model_data["pipeline"] caption = pipeline(image)[0]['generated_text'] return JSONResponse({"caption": caption}) elif request.task_type == "audio-transcription": if request.audio_file is None: raise HTTPException(status_code=400, detail="Audio file is required for audio transcription.") try: contents = await request.audio_file.read() pipeline = model_data["pipeline"] try: transcription = pipeline(contents, sampling_rate=16000)[0]["text"] # Assuming 16kHz sampling rate return JSONResponse({"transcription": transcription}) except Exception as e: raise HTTPException(status_code=500, detail=f"Error during audio transcription (pipeline): {str(e)}") except Exception as e: raise HTTPException(status_code=500, detail=f"Error during audio transcription (file read): {str(e)}") elif request.task_type == "summarization": if request.input_text is None: raise HTTPException(status_code=400, detail="Input text is required for summarization.") model = model_data["model"].to(device) tokenizer = model_data["tokenizer"] inputs = tokenizer(request.input_text, return_tensors="pt", truncation=True, max_length=512) # added max_length for summarization with torch.no_grad(): outputs = model.generate(**inputs) summary = tokenizer.decode(outputs[0], skip_special_tokens=True) return JSONResponse({"summary": summary}) else: raise HTTPException(status_code=500, detail=f"Unsupported task type") except Exception as e: logger.exception(f"Internal server error: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") @app.get("/", response_class=HTMLResponse) async def root(request: Request): return TEMPLATES.TemplateResponse("index.html", {"request": request}) @app.get("/health") async def health_check(): return {"status": "healthy"} # Authentication Endpoints @app.post("/token", response_model=Token) async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()): user = authenticate_user(form_data.username, form_data.password) if not user: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect username or password", headers={"WWW-Authenticate": "Bearer"}, ) access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES) access_token = create_access_token(data={"sub": user["username"]}, expires_delta=access_token_expires) return {"access_token": access_token, "token_type": "bearer"} def authenticate_user(username: str, password: str): user = get_user(username) if user and pwd_context.verify(password, user.hashed_password): return {"username": user.username} return None def create_access_token(data: Dict[str, Any], expires_delta: timedelta = None): to_encode = data.copy() if expires_delta: expire = datetime.utcnow() + expires_delta else: expire = datetime.utcnow() + timedelta(minutes=15) to_encode.update({"exp": expire}) encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) return encoded_jwt class Token(BaseModel): access_token: str token_type: str @app.get("/users/me") async def read_users_me(current_user: str = Depends(get_current_user)): return {"username": current_user} async def get_current_user(token: str = Depends(oauth2_scheme)): credentials_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"}, ) try: payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) username: str = payload.get("sub") if username is None: raise credentials_exception token_data = {"username": username, "token": token} except JWTError: raise credentials_exception user = get_user(username) if user is None: raise credentials_exception return username @app.post("/register", response_model=User, status_code=status.HTTP_201_CREATED) async def create_user(user: User): try: hashed_password = pwd_context.hash(user.password) new_user = {"username": user.username, "email": user.email, "hashed_password": hashed_password} inserted_user = insert_user(new_user) if inserted_user: return User(**inserted_user) else: raise HTTPException(status_code=500, detail="Failed to create user.") except Exception as e: logger.error(f"Error creating user: {e}") raise HTTPException(status_code=500, detail=f"Error creating user: {e}") @app.put("/users/{username}", response_model=User, dependencies=[Depends(get_current_user)]) async def update_user_data(username: str, user: User): try: hashed_password = pwd_context.hash(user.password) updated_user_data = {"email": user.email, "hashed_password": hashed_password} updated_user = update_user(username, updated_user_data) if updated_user: return User(**updated_user) else: raise HTTPException(status_code=404, detail="User not found") except Exception as e: logger.error(f"Error updating user: {e}") raise HTTPException(status_code=500, detail="Error updating user.") @app.delete("/users/{username}", dependencies=[Depends(get_current_user)]) async def delete_user_account(username: str): try: deleted_user = delete_user(username) if deleted_user: return JSONResponse({"message": "User deleted successfully."}, status_code=200) else: raise HTTPException(status_code=404, detail="User not found") except Exception as e: logger.error(f"Error deleting user: {e}") raise HTTPException(status_code=500, detail="Error deleting user.") @app.get("/users", dependencies=[Depends(get_current_user)]) async def get_all_users_route(): return get_all_users() @app.exception_handler(RequestValidationError) async def validation_exception_handler(request: Request, exc: RequestValidationError): return JSONResponse( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, content=json.dumps({"detail": exc.errors(), "body": exc.body}), ) if __name__ == "__main__": create_db_and_table() # Initialize database on startup uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True) # replace main with your filename