Upload app.py

src/apps/app.py

@@ -16,7 +16,7 @@ from sqlalchemy.future import select
 from pydantic import BaseModel, EmailStr
 
 from utils.main import RAG
-from database import engine, Base, get_db, AsyncSessionLocal
+from database import engine, Base, get_db
 from models import User
 from auth import get_password_hash, verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM, SECRET_KEY
 from email_utils import send_reset_email
@@ -25,6 +25,8 @@ from fastapi.security import OAuth2PasswordBearer
 from jose import JWTError, jwt
 import uuid
 
+from database import engine, Base, get_db, AsyncSessionLocal
+
 # Initialize Database and Create Default Admin
 async def init_db():
     async with engine.begin() as conn:
@@ -338,11 +340,11 @@ async def reset_password_page(request: Request):
 # Home and Role Selection
 @app.get("/", response_class=HTMLResponse)
 async def role_selection(request: Request):
-    return templates.TemplateResponse("roleselection.html", {"request": request})
+    return templates.TemplateResponse("roleselection.html", {"request": request, "show_roles": False})
 
 @app.get("/role", response_class=HTMLResponse)
 async def roleselection_page(request: Request):
-    return templates.TemplateResponse("roleselection.html", {"request": request})
+    return templates.TemplateResponse("roleselection.html", {"request": request, "show_roles": True})
 
 # Chatbot Pages
 @app.get("/judgechatbot.html", response_class=HTMLResponse)
@@ -395,6 +397,8 @@ async def advocatechatbot_page(request: Request):
 async def womanchatbot_page(request: Request):
     return templates.TemplateResponse("womanchatbot.html", {"request": request})
 
+
+
 @app.get("/safetytips.html", response_class=HTMLResponse)
 async def safetytips_page(request: Request):
     return templates.TemplateResponse("safetytips.html", {"request": request})
@@ -420,16 +424,23 @@ async def stream_text_conversational(websocket: WebSocket, query: str, role: str
     chat_limit = 10
     temp_chat = {"user": "" ,"system":""}
     temp_chat["user"] = query
+    # print(f"DEBUG: stream_text_conversational executing with role={role}")
     model_response = ""
     try:
-        completion = RAG(query, chat_history, role=role)
+        if role == "Citizen":
+            completion = RAG(query, chat_history, role=role) 
+        else:
+            completion = RAG(query, chat_history, role=role)
         for chunk in completion:
             if chunk.choices[0].delta.content is not None:
                 await websocket.send_text(chunk.choices[0].delta.content)
                 await asyncio.sleep(0.01)
                 model_response += chunk.choices[0].delta.content
         
+        # Signal completion to frontend
         await websocket.send_text("[DONE]")
+        
+        # print(model_response)
         temp_chat['system']=model_response
         chat_history.append(temp_chat)
         if len(chat_history)>chat_limit:
@@ -444,6 +455,8 @@ async def stream_text_conversational(websocket: WebSocket, query: str, role: str
 async def conversational_chat(websocket: WebSocket, role: Optional[str] = None):
     await websocket.accept()
     
+    # --- STRICT ROLE ENFORCEMENT ---
+    # 1. Get role from Token/Cookie (Source of Truth)
     token = websocket.cookies.get("access_token")
     token_role = "General"
     user_id = None
@@ -470,6 +483,7 @@ async def conversational_chat(websocket: WebSocket, role: Optional[str] = None):
         await websocket.close()
         return
     
+    # 2. Apply Logic based on Source of Truth
     if token_role == "Admin":
         if not role:
             role = "Admin"
@@ -480,6 +494,7 @@ async def conversational_chat(websocket: WebSocket, role: Optional[str] = None):
         try:
             query = await websocket.receive_text()
             
+            # --- USAGE LIMIT CHECK ---
             if token_role != "Admin" and question_count >= 2:
                 limit_message = (
                     "### Free usage limit reached\n\n"
@@ -494,13 +509,14 @@ async def conversational_chat(websocket: WebSocket, role: Optional[str] = None):
             print(f"Query ({role}): {query}")
             await stream_text_conversational(websocket, query, role=role)
             
+            # --- INCREMENT USAGE ---
             if token_role != "Admin":
                 async for db in get_db():
                     result = await db.execute(select(User).where(User.id == user_id))
                     user_to_update = result.scalars().first()
                     if user_to_update:
                         user_to_update.question_count += 1
-                        question_count = user_to_update.question_count
+                        question_count = user_to_update.question_count # Keep local sync
                         await db.commit()
                     break
                     
@@ -515,7 +531,9 @@ async def conversational_chat(websocket: WebSocket, role: Optional[str] = None):
                 pass
             break
 
+
 # Run the application
 if __name__ == "__main__":
+    print("Starting Law Bot Server...")
     port = int(os.getenv("PORT", 8000))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    uvicorn.run(app, host="0.0.0.0", port=port)