bearer-auth

docs/menv.md

@@ -4,6 +4,6 @@ export SPEECH_RATE_LIMIT=5/minute
 export CHAT_RATE_LIMIT=100/minute
 export EXTERNAL_TTS_URL=https://slabstech-dhwani-internal-api-server.hf.space/v1/audio/speech
 export EXTERNAL_ASR_URL=https://gaganyatri-asr-indic-server-cpu.hf.space
-export EXTERNAL_TEXT_GEN_URL=https://gaganyatri-asr-indic-server-cpu.hf.space
-export EXTERNAL_AUDIO_PROC_URL=https://gaganyatri-asr-indic-server-cpu.hf.space
+export EXTERNAL_TEXT_GEN_URL=https://slabstech-dhwani-internal-api-server.hf.space
+export EXTERNAL_AUDIO_PROC_URL=https://slabstech-dhwani-internal-api-server.hf.space
 export API_KEY_SECRET=your_secret_key

requirements.txt

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

src/server/main.py

@@ -9,40 +9,17 @@ from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, Uploa
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
 from pydantic import BaseModel, Field, field_validator
-from pydantic_settings import BaseSettings
 from slowapi import Limiter
 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
+
 # Assuming these are in your project structure
 from config.tts_config import SPEED, ResponseFormat, config as tts_config
 from config.logging_config import logger
-#from utils.auth import get_api_key
-
-# Configuration settings
-class Settings(BaseSettings):
-    llm_model_name: str = "google/gemma-3-4b-it"
-    max_tokens: int = 512
-    host: str = "0.0.0.0"
-    port: int = 7860
-    chat_rate_limit: str = "100/minute"
-    speech_rate_limit: str = "5/minute"
-    external_tts_url: str = Field(..., env="EXTERNAL_TTS_URL")
-    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")
-    api_key_secret: str = Field(..., env="API_KEY_SECRET")
-
-    @field_validator("chat_rate_limit", "speech_rate_limit")
-    def validate_rate_limit(cls, v):
-        if not v.count("/") == 1 or not v.split("/")[0].isdigit():
-            raise ValueError("Rate limit must be in format 'number/period' (e.g., '5/minute')")
-        return v
-
-    class Config:
-        env_file = ".env"
-        env_file_encoding = "utf-8"
 
 settings = Settings()
 
@@ -127,12 +104,17 @@ async def health_check():
 async def home():
     return RedirectResponse(url="/docs")
 
+@app.post("/v1/token", response_model=TokenResponse)
+async def token(user_id: str = Form(...)):
+    # In production, add proper authentication (e.g., password validation)
+    return await login(user_id=user_id)
+
 @app.post("/v1/audio/speech")
 @limiter.limit(settings.speech_rate_limit)
 async def generate_audio(
     request: Request,
     speech_request: SpeechRequest = Depends(),
-    #api_key: str = Depends(get_api_key),
+    user_id: str = Depends(get_current_user),
     tts_service: TTSService = Depends(get_tts_service)
 ):
     if not speech_request.input.strip():
@@ -141,7 +123,8 @@ async def generate_audio(
     logger.info("Processing speech request", extra={
         "endpoint": "/v1/audio/speech",
         "input_length": len(speech_request.input),
-        "client_ip": get_remote_address(request)
+        "client_ip": get_remote_address(request),
+        "user_id": user_id
     })
     
     payload = {
@@ -167,10 +150,9 @@ async def generate_audio(
         headers=headers
     )
 
-
 class ChatRequest(BaseModel):
     prompt: str
-    src_lang: str = "kan_Knda"  # Default to Kannada
+    src_lang: str = "kan_Knda"
 
     @field_validator("prompt")
     def prompt_must_be_valid(cls, v):
@@ -181,22 +163,23 @@ class ChatRequest(BaseModel):
 class ChatResponse(BaseModel):
     response: str
 
-
 @app.post("/v1/chat", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
-async def chat(request: Request, chat_request: ChatRequest):
+async def chat(
+    request: Request,
+    chat_request: ChatRequest,
+    user_id: str = Depends(get_current_user)
+):
     if not chat_request.prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}")
+    logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, user_id: {user_id}")
     
     try:
-
-        # Call the external API instead of llm_manager.generate
         external_url = "https://slabstech-dhwani-internal-api-server.hf.space/v1/chat"
         payload = {
-            "prompt": chat_request.prompt ,
-            "src_lang": chat_request.src_lang,  
-            "tgt_lang" : chat_request.src_lang
+            "prompt": chat_request.prompt,
+            "src_lang": chat_request.src_lang,
+            "tgt_lang": chat_request.src_lang
         }
         
         response = requests.post(
@@ -208,14 +191,12 @@ async def chat(request: Request, chat_request: ChatRequest):
             },
             timeout=60
         )
-        response.raise_for_status()  # Raise an exception for bad status codes
+        response.raise_for_status()
         
-        # Extract the response text from the API
         response_data = response.json()
-        response = response_data.get("response", "")
-        logger.info(f"Generated Chat response from external API: {response}")
-
-        return ChatResponse(response=response)
+        response_text = response_data.get("response", "")
+        logger.info(f"Generated Chat response from external API: {response_text}")
+        return ChatResponse(response=response_text)
     
     except requests.Timeout:
         logger.error("External chat API request timed out")
@@ -232,13 +213,14 @@ async def chat(request: Request, chat_request: ChatRequest):
 async def process_audio(
     file: UploadFile = File(...),
     language: str = Query(..., enum=["kannada", "hindi", "tamil"]),
-    #api_key: str = Depends(get_api_key),
+    user_id: str = Depends(get_current_user),
     request: Request = None,
 ):
     logger.info("Processing audio processing request", extra={
         "endpoint": "/v1/process_audio",
         "filename": file.filename,
-        "client_ip": get_remote_address(request)
+        "client_ip": get_remote_address(request),
+        "user_id": user_id
     })
     
     start_time = time()
@@ -269,16 +251,9 @@ async def process_audio(
 async def transcribe_audio(
     file: UploadFile = File(...),
     language: str = Query(..., enum=["kannada", "hindi", "tamil"]),
-    #api_key: str = Depends(get_api_key),
+    user_id: str = Depends(get_current_user),
     request: Request = None,
 ):
-    '''
-    logger.info("Processing transcription request", extra={
-        "endpoint": "/v1/transcribe",
-        "filename": file.filename,
-        "client_ip": get_remote_address(request)
-    })
-    '''
     start_time = time()
     try:
         file_content = await file.read()
@@ -294,13 +269,11 @@ async def transcribe_audio(
         response.raise_for_status()
         
         transcription = response.json().get("text", "")
-        #logger.info(f"Transcription completed in {time() - start_time:.2f} seconds")
         return TranscriptionResponse(text=transcription)
     
     except requests.Timeout:
         raise HTTPException(status_code=504, detail="Transcription service timeout")
     except requests.RequestException as e:
-        #logger.error(f"Transcription request failed: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
 
 @app.post("/v1/chat_v2", response_model=TranscriptionResponse)
@@ -309,7 +282,7 @@ async def chat_v2(
     request: Request,
     prompt: str = Form(...),
     image: UploadFile = File(default=None),
-    #api_key: str = Depends(get_api_key)
+    user_id: str = Depends(get_current_user)
 ):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
@@ -318,18 +291,18 @@ async def chat_v2(
         "endpoint": "/v1/chat_v2",
         "prompt_length": len(prompt),
         "has_image": bool(image),
-        "client_ip": get_remote_address(request)
+        "client_ip": get_remote_address(request),
+        "user_id": user_id
     })
     
     try:
-        # For demonstration, we'll just return the prompt as text
         image_data = Image.open(await image.read()) if image else None
         response_text = f"Processed: {prompt}" + (" with image" if image_data else "")
         return TranscriptionResponse(text=response_text)
     except Exception as e:
         logger.error(f"Chat_v2 processing failed: {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
-    
+
 class TranslationRequest(BaseModel):
     sentences: list[str]
     src_lang: str
@@ -339,13 +312,14 @@ class TranslationResponse(BaseModel):
     translations: list[str]
 
 @app.post("/v1/translate", response_model=TranslationResponse)
-async def translate(request: TranslationRequest):
-    logger.info(f"Received translation request: {request.dict()}")
+async def translate(
+    request: TranslationRequest,
+    user_id: str = Depends(get_current_user)
+):
+    logger.info(f"Received translation request: {request.dict()}, user_id: {user_id}")
     
-    # External API endpoint
     external_url = f"https://slabstech-dhwani-internal-api-server.hf.space/translate?src_lang={request.src_lang}&tgt_lang={request.tgt_lang}"
     
-    # Prepare the payload matching the external API's expected format
     payload = {
         "sentences": request.sentences,
         "src_lang": request.src_lang,
@@ -353,7 +327,6 @@ async def translate(request: TranslationRequest):
     }
     
     try:
-        # Make the POST request to the external API
         response = requests.post(
             external_url,
             json=payload,
@@ -361,13 +334,10 @@ async def translate(request: TranslationRequest):
                 "accept": "application/json",
                 "Content-Type": "application/json"
             },
-            timeout=60  # Set a timeout to avoid hanging
+            timeout=60
         )
-        
-        # Raise an exception for bad status codes (4xx, 5xx)
         response.raise_for_status()
         
-        # Extract translations from the response
         response_data = response.json()
         translations = response_data.get("translations", [])
         
@@ -388,8 +358,6 @@ async def translate(request: TranslationRequest):
         logger.error(f"Invalid JSON response: {str(e)}")
         raise HTTPException(status_code=500, detail="Invalid response format from translation service")
 
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run the FastAPI server.")
     parser.add_argument("--port", type=int, default=settings.port, help="Port to run the server on.")

src/server/utils/auth.py

@@ -1,21 +1,94 @@
-from fastapi.security import APIKeyHeader
+import jwt
+from datetime import datetime, timedelta
+from pydantic import BaseModel, Field, field_validator
+
+from fastapi.security import OAuth2PasswordBearer
 from fastapi import HTTPException, status, Depends
+from pydantic import BaseModel
 from pydantic_settings import BaseSettings
-from config.logging_config import logger
+from config.logging_config import logger  # Assuming this is available
+from typing import Optional
 
+# Centralized Settings class (can be moved to a separate config file later)
 class Settings(BaseSettings):
-    api_key: str
+    api_key_secret: str = Field(..., env="API_KEY_SECRET")  # Secret key for signing JWTs
+    token_expiration_minutes: int = Field(30, env="TOKEN_EXPIRATION_MINUTES")  # Default to 30 minutes
+    llm_model_name: str = "google/gemma-3-4b-it"
+    max_tokens: int = 512
+    host: str = "0.0.0.0"
+    port: int = 7860
+    chat_rate_limit: str = "100/minute"
+    speech_rate_limit: str = "5/minute"
+
     class Config:
         env_file = ".env"
+        env_file_encoding = "utf-8"
 
 settings = Settings()
+logger.info(f"Loaded API_KEY_SECRET at startup: {settings.api_key_secret}")  # Add this line
+
+# OAuth2 scheme with Bearer token
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/v1/token")
+
+# Model for token payload
+class TokenPayload(BaseModel):
+    sub: str  # Subject (user identifier)
+    exp: int  # Expiration timestamp
+
+# Model for token response
+class TokenResponse(BaseModel):
+    access_token: str
+    token_type: str
+
+async def create_access_token(user_id: str) -> str:
+    """
+    Create a JWT access token for a given user.
+    """
+    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}")  # Add this line
+    token = jwt.encode(payload, settings.api_key_secret, algorithm="HS256")
+    logger.info(f"Generated access token for user: {user_id}")
+    return token
 
-API_KEY_NAME = "X-API-Key"
-api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+async def get_current_user(token: str = Depends(oauth2_scheme)) -> str:
+    """
+    Validate the Bearer token and return the user ID.
+    """
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid authentication credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    
+    try:
+        logger.info(f"Verifying token with API_KEY_SECRET: {settings.api_key_secret}")  # Add this line
+        payload = jwt.decode(token, settings.api_key_secret, algorithms=["HS256"])
+        token_data = TokenPayload(**payload)
+        user_id = token_data.sub
+        if user_id is None:
+            raise credentials_exception
+        if datetime.utcnow().timestamp() > token_data.exp:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Token has expired",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        logger.info(f"Validated token for user: {user_id}")
+        return user_id
+    except jwt.InvalidTokenError:
+        logger.warning(f"Invalid token attempt: {token[:10]}...")
+        raise credentials_exception
+    except Exception as e:
+        logger.error(f"Token validation error: {str(e)}")
+        raise credentials_exception
 
-async def get_api_key(api_key: str = Depends(api_key_header)):
-    if api_key != settings.api_key:
-        logger.warning(f"Failed API key attempt: {api_key}")
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key")
-    logger.info("API key validated successfully")
-    return api_key
+# For demonstration purposes, a simple login function
+# In production, replace with proper user authentication (e.g., database lookup)
+async def login(user_id: str) -> TokenResponse:
+    """
+    Generate a token for a user. In production, validate credentials here.
+    """
+    # Placeholder: Assume user_id is valid; in reality, check against a database
+    token = await create_access_token(user_id=user_id)
+    return TokenResponse(access_token=token, token_type="bearer")