update-refresh token

src/server/main.py

@@ -5,7 +5,7 @@ from typing import List, Optional
 from abc import ABC, abstractmethod
 
 import uvicorn
-from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Form, Security
+from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
@@ -147,7 +147,7 @@ async def home():
 @app.post("/v1/token", 
           response_model=TokenResponse,
           summary="User Login",
-          description="Authenticate a user with username and password to obtain an access token. Copy the token and use it in the 'Authorize' button above.",
+          description="Authenticate a user with username and password to obtain an access token and refresh token. Copy the access token and use it in the 'Authorize' button above.",
           tags=["Authentication"],
           responses={
               200: {"description": "Successful login", "model": TokenResponse},
@@ -159,11 +159,11 @@ async def token(login_request: LoginRequest):
 @app.post("/v1/refresh", 
           response_model=TokenResponse,
           summary="Refresh Access Token",
-          description="Generate a new access token using an existing valid token.",
+          description="Generate a new access token and refresh token using an existing valid refresh token.",
           tags=["Authentication"],
           responses={
-              200: {"description": "New token issued", "model": TokenResponse},
-              401: {"description": "Invalid or expired token"}
+              200: {"description": "New tokens issued", "model": TokenResponse},
+              401: {"description": "Invalid or expired refresh token"}
           })
 async def refresh(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)):
     return await refresh_token(credentials)
@@ -171,7 +171,7 @@ async def refresh(credentials: HTTPAuthorizationCredentials = Depends(bearer_sch
 @app.post("/v1/register", 
           response_model=TokenResponse,
           summary="Register New User",
-          description="Create a new user account and return an access token. Requires admin access (use 'admin' user with password 'adminpass' initially).",
+          description="Create a new user account and return an access token and refresh token. Requires admin access (use 'admin' user with password 'adminpass' initially).",
           tags=["Authentication"],
           responses={
               200: {"description": "User registered successfully", "model": TokenResponse},
@@ -180,9 +180,9 @@ async def refresh(credentials: HTTPAuthorizationCredentials = Depends(bearer_sch
           })
 async def register_user(
     register_request: RegisterRequest,
-    current_user: str = Depends(get_current_user_with_admin)  # Enforce admin-only access
+    current_user: str = Depends(get_current_user_with_admin)
 ):
-    return await register(register_request, current_user)  # Pass current_user explicitly
+    return await register(register_request, current_user)
 
 @app.post("/v1/audio/speech",
           summary="Generate Speech from Text",
@@ -520,9 +520,7 @@ async def translate(
     except ValueError as e:
         logger.error(f"Invalid JSON response: {str(e)}")
         raise HTTPException(status_code=500, detail="Invalid response format from translation service")
-    
 
-# Request/Response Models for Visual Query
 class VisualQueryRequest(BaseModel):
     query: str
     src_lang: str = "kan_Knda"  # Default to Kannada
@@ -537,49 +535,48 @@ class VisualQueryRequest(BaseModel):
 class VisualQueryResponse(BaseModel):
     answer: str
 
-
-@app.post("/v1/visual_query", response_model=VisualQueryResponse)
+@app.post("/v1/visual_query", 
+          response_model=VisualQueryResponse,
+          summary="Visual Query with Image",
+          description="Process a visual query with an image and text question. Rate limited to 100 requests per minute per user. Requires authentication.",
+          tags=["Chat"],
+          responses={
+              200: {"description": "Query response", "model": VisualQueryResponse},
+              400: {"description": "Invalid query"},
+              401: {"description": "Unauthorized - Token required"},
+              429: {"description": "Rate limit exceeded"},
+              504: {"description": "Visual query service timeout"}
+          })
 @limiter.limit(settings.chat_rate_limit)
 async def visual_query(
     request: Request,
-    query: str = Form(...),
-    file: UploadFile = File(...),
-    src_lang: str = Query(default="kan_Knda"),
-    tgt_lang: str = Query(default="kan_Knda"),
-    #api_key: str = Depends(get_api_key)  # Uncomment if authentication is needed
+    query: str = Form(..., description="Text query for the visual content"),
+    file: UploadFile = File(..., description="Image file to analyze"),
+    src_lang: str = Query(default="kan_Knda", description="Source language code"),
+    tgt_lang: str = Query(default="kan_Knda", description="Target language code"),
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
 ):
-    """
-    Endpoint to process visual queries with an image and text question.
-    Calls an external API to analyze the image and provide a response.
-    """
+    user_id = await get_current_user(credentials)
     if not query.strip():
         raise HTTPException(status_code=400, detail="Query cannot be empty")
     
     logger.info("Processing visual query request", extra={
         "endpoint": "/v1/visual_query",
         "query_length": len(query),
-        "file_name": file.filename,  # Changed from "filename" to "file_name"
+        "file_name": file.filename,
         "client_ip": get_remote_address(request),
+        "user_id": user_id,
         "src_lang": src_lang,
         "tgt_lang": tgt_lang
     })
     
-    # External API URL
     external_url = f"https://slabstech-dhwani-internal-api-server.hf.space/v1/visual_query/?src_lang={src_lang}&tgt_lang={tgt_lang}"
     
     try:
-        # Read file content
         file_content = await file.read()
+        files = {"file": (file.filename, file_content, file.content_type)}
+        data = {"query": query}
         
-        # Prepare multipart/form-data
-        files = {
-            "file": (file.filename, file_content, file.content_type)
-        }
-        data = {
-            "query": query
-        }
-        
-        # Make the POST request to external API
         response = requests.post(
             external_url,
             files=files,
@@ -587,11 +584,8 @@ async def visual_query(
             headers={"accept": "application/json"},
             timeout=60
         )
-        
-        # Raise an exception for bad status codes
         response.raise_for_status()
         
-        # Parse the response
         response_data = response.json()
         answer = response_data.get("answer", "")
         
@@ -611,9 +605,6 @@ async def visual_query(
     except ValueError as e:
         logger.error(f"Invalid JSON response: {str(e)}")
         raise HTTPException(status_code=500, detail="Invalid response format from visual query service")
-    except Exception as e:
-        logger.error(f"Unexpected error in visual query: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run the FastAPI server.")

src/server/utils/auth.py

@@ -29,7 +29,8 @@ pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 class Settings(BaseSettings):
     api_key_secret: str = Field(..., env="API_KEY_SECRET")
-    token_expiration_minutes: int = Field(30, env="TOKEN_EXPIRATION_MINUTES")
+    token_expiration_minutes: int = Field(1440, env="TOKEN_EXPIRATION_MINUTES")  # 24 hours
+    refresh_token_expiration_days: int = Field(7, env="REFRESH_TOKEN_EXPIRATION_DAYS")  # 7 days
     llm_model_name: str = "google/gemma-3-4b-it"
     max_tokens: int = 512
     host: str = "0.0.0.0"
@@ -40,7 +41,6 @@ class Settings(BaseSettings):
     external_asr_url: str = Field(..., env="EXTERNAL_ASR_URL")
     external_text_gen_url: str = Field(..., env="EXTERNAL_TEXT_GEN_URL")
     external_audio_proc_url: str = Field(..., env="EXTERNAL_AUDIO_PROC_URL")
-    # Admin credentials required from environment variables, no defaults
     default_admin_username: str = Field("admin", env="DEFAULT_ADMIN_USERNAME")
     default_admin_password: str = Field("admin54321", env="DEFAULT_ADMIN_PASSWORD")
 
@@ -49,7 +49,6 @@ class Settings(BaseSettings):
         env_file_encoding = "utf-8"
 
 settings = Settings()
-#logger.info(f"Loaded API_KEY_SECRET at startup: {settings.api_key_secret}")
 
 # Seed initial data (optional)
 def seed_initial_data():
@@ -77,9 +76,11 @@ bearer_scheme = HTTPBearer()
 class TokenPayload(BaseModel):
     sub: str
     exp: float
+    type: str
 
 class TokenResponse(BaseModel):
     access_token: str
+    refresh_token: str
     token_type: str
 
 class LoginRequest(BaseModel):
@@ -90,13 +91,15 @@ class RegisterRequest(BaseModel):
     username: str
     password: str
 
-async def create_access_token(user_id: str) -> str:
+async def create_access_token(user_id: str) -> dict:
     expire = datetime.utcnow() + timedelta(minutes=settings.token_expiration_minutes)
-    payload = {"sub": user_id, "exp": expire.timestamp()}
-    #logger.info(f"Signing token with API_KEY_SECRET: {settings.api_key_secret}")
+    payload = {"sub": user_id, "exp": expire.timestamp(), "type": "access"}
     token = jwt.encode(payload, settings.api_key_secret, algorithm="HS256")
-    #logger.info(f"Generated access token for user: {user_id}")
-    return token
+    refresh_expire = datetime.utcnow() + timedelta(days=settings.refresh_token_expiration_days)
+    refresh_payload = {"sub": user_id, "exp": refresh_expire.timestamp(), "type": "refresh"}
+    refresh_token = jwt.encode(refresh_payload, settings.api_key_secret, algorithm="HS256")
+    logger.info(f"Generated tokens for user: {user_id}")
+    return {"access_token": token, "refresh_token": refresh_token}
 
 async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> str:
     token = credentials.credentials
@@ -106,10 +109,7 @@ async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(b
         headers={"WWW-Authenticate": "Bearer"},
     )
     try:
-        #logger.info(f"Received token: {token}")
-        #logger.info(f"Verifying token with API_KEY_SECRET: {settings.api_key_secret}")
         payload = jwt.decode(token, settings.api_key_secret, algorithms=["HS256"], options={"verify_exp": False})
-        #logger.info(f"Decoded payload: {payload}")
         token_data = TokenPayload(**payload)
         user_id = token_data.sub
         
@@ -121,9 +121,8 @@ async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(b
             raise credentials_exception
         
         current_time = datetime.utcnow().timestamp()
-        logger.info(f"Current time: {current_time}, Token exp: {token_data.exp}")
         if current_time > token_data.exp:
-            #logger.warning(f"Token expired: current_time={current_time}, exp={token_data.exp}")
+            logger.warning(f"Token expired: current_time={current_time}, exp={token_data.exp}")
             raise HTTPException(
                 status_code=status.HTTP_401_UNAUTHORIZED,
                 detail="Token has expired",
@@ -159,8 +158,8 @@ async def login(login_request: LoginRequest) -> TokenResponse:
     if not user or not pwd_context.verify(login_request.password, user.password):
         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=user.username)
-    return TokenResponse(access_token=token, token_type="bearer")
+    tokens = await create_access_token(user_id=user.username)
+    return TokenResponse(access_token=tokens["access_token"], refresh_token=tokens["refresh_token"], token_type="bearer")
 
 async def register(register_request: RegisterRequest, current_user: str = Depends(get_current_user_with_admin)) -> TokenResponse:
     db = SessionLocal()
@@ -176,11 +175,26 @@ async def register(register_request: RegisterRequest, current_user: str = Depend
     db.commit()
     db.close()
     
-    token = await create_access_token(user_id=register_request.username)
+    tokens = await create_access_token(user_id=register_request.username)
     logger.info(f"Registered and generated token for user: {register_request.username} by admin {current_user}")
-    return TokenResponse(access_token=token, token_type="bearer")
+    return TokenResponse(access_token=tokens["access_token"], refresh_token=tokens["refresh_token"], token_type="bearer")
 
 async def refresh_token(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> TokenResponse:
-    user_id = await get_current_user(credentials)
-    new_token = await create_access_token(user_id=user_id)
-    return TokenResponse(access_token=new_token, token_type="bearer")
+    token = credentials.credentials
+    try:
+        payload = jwt.decode(token, settings.api_key_secret, algorithms=["HS256"])
+        token_data = TokenPayload(**payload)
+        if payload.get("type") != "refresh":
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token type; refresh token required")
+        user_id = token_data.sub
+        db = SessionLocal()
+        user = db.query(User).filter_by(username=user_id).first()
+        db.close()
+        if not user:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found")
+        if datetime.utcnow().timestamp() > token_data.exp:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Refresh token has expired")
+        tokens = await create_access_token(user_id=user_id)
+        return TokenResponse(access_token=tokens["access_token"], refresh_token=tokens["refresh_token"], token_type="bearer")
+    except jwt.InvalidTokenError:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid refresh token")