add-user/

requirements.txt

@@ -5,4 +5,5 @@ slowapi
 requests
 python-multipart
 pillow
-pyjwt
+pyjwt
+sqlalchemy

src/server/main.py

@@ -14,9 +14,7 @@ from slowapi.util import get_remote_address
 import requests
 from PIL import Image
 
-# Import from auth.py
-#from utils.auth import get_current_user, login, TokenResponse, Settings
-from utils.auth import get_current_user, login, TokenResponse, Settings, LoginRequest
+from utils.auth import get_current_user, login, refresh_token, TokenResponse, Settings, LoginRequest
 
 # Assuming these are in your project structure
 from config.tts_config import SPEED, ResponseFormat, config as tts_config
@@ -96,7 +94,16 @@ class ExternalTTSService(TTSService):
 def get_tts_service() -> TTSService:
     return ExternalTTSService()
 
-# Endpoints
+
+
+@app.post("/v1/token", response_model=TokenResponse)
+async def token(login_request: LoginRequest):
+    return await login(login_request)
+
+@app.post("/v1/refresh", response_model=TokenResponse)
+async def refresh(token_response: TokenResponse = Depends(refresh_token)):
+    return token_response
+
 @app.get("/v1/health")
 async def health_check():
     return {"status": "healthy", "model": settings.llm_model_name}
@@ -105,9 +112,6 @@ async def health_check():
 async def home():
     return RedirectResponse(url="/docs")
 
-@app.post("/v1/token", response_model=TokenResponse)
-async def token(login_request: LoginRequest):
-    return await login(login_request)
 
 @app.post("/v1/audio/speech")
 @limiter.limit(settings.speech_rate_limit)

src/server/utils/auth.py

@@ -5,7 +5,33 @@ from fastapi import HTTPException, status, Depends
 from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings
 from config.logging_config import logger
-from typing import Dict
+from sqlalchemy import create_engine, Column, String
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+# SQLite database setup
+DATABASE_URL = "sqlite:///users.db"
+engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})  # For SQLite threading
+Base = declarative_base()
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+class User(Base):
+    __tablename__ = "users"
+    username = Column(String, primary_key=True, index=True)
+    password = Column(String)  # In production, hash this with bcrypt
+
+# Create the database tables
+Base.metadata.create_all(bind=engine)
+
+# Seed initial data (optional, for testing)
+def seed_initial_data():
+    db = SessionLocal()
+    if not db.query(User).filter_by(username="testuser").first():
+        db.add(User(username="testuser", password="password123"))
+        db.commit()
+    db.close()
+
+seed_initial_data()  # Run once at startup
 
 class Settings(BaseSettings):
     api_key_secret: str = Field(..., env="API_KEY_SECRET")
@@ -38,13 +64,6 @@ class TokenResponse(BaseModel):
     access_token: str
     token_type: str
 
-# Simple in-memory user store (replace with database in production)
-# Format: {username: password}
-USERS_DB: Dict[str, str] = {
-    "testuser": "password123",
-    "admin": "adminpass"
-}
-
 class LoginRequest(BaseModel):
     username: str
     password: str
@@ -71,7 +90,10 @@ async def get_current_user(token: str = Depends(oauth2_scheme)) -> str:
         token_data = TokenPayload(**payload)
         user_id = token_data.sub
         
-        if user_id is None or user_id not in USERS_DB:
+        db = SessionLocal()
+        user = db.query(User).filter_by(username=user_id).first()
+        db.close()
+        if user_id is None or not user:
             logger.warning(f"Invalid or unknown user: {user_id}")
             raise credentials_exception
         
@@ -98,10 +120,16 @@ async def get_current_user(token: str = Depends(oauth2_scheme)) -> str:
         raise credentials_exception
 
 async def login(login_request: LoginRequest) -> TokenResponse:
-    username = login_request.username
-    password = login_request.password
-    if username not in USERS_DB or USERS_DB[username] != password:
-        logger.warning(f"Login failed for user: {username}")
+    db = SessionLocal()
+    user = db.query(User).filter_by(username=login_request.username).first()
+    db.close()
+    if not user or user.password != login_request.password:  # Use password hashing in production
+        logger.warning(f"Login failed for user: {login_request.username}")
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid username or password")
-    token = await create_access_token(user_id=username)
-    return TokenResponse(access_token=token, token_type="bearer")
+    token = await create_access_token(user_id=user.username)
+    return TokenResponse(access_token=token, token_type="bearer")
+
+async def refresh_token(token: str = Depends(oauth2_scheme)) -> TokenResponse:
+    user_id = await get_current_user(token)
+    new_token = await create_access_token(user_id=user_id)
+    return TokenResponse(access_token=new_token, token_type="bearer")