Update src/apps/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
+from database import engine, Base, get_db, AsyncSessionLocal
 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,10 +25,27 @@ from fastapi.security import OAuth2PasswordBearer
 from jose import JWTError, jwt
 import uuid
 
-# Initialize Database
+# Initialize Database and Create Default Admin
 async def init_db():
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
+    
+    # Ensure Admin exists
+    async with AsyncSessionLocal() as db:
+        result = await db.execute(select(User).where(User.username == "admin"))
+        admin_user = result.scalars().first()
+        
+        if not admin_user:
+            print("Creating default admin user...")
+            new_admin = User(
+                username="admin",
+                email="admin@lawbot.ai",
+                hashed_password=get_password_hash("ADMin"),
+                role="Admin"
+            )
+            db.add(new_admin)
+            await db.commit()
+            print("Admin user created successfully.")
 
 app = FastAPI(on_startup=[init_db])
 
@@ -378,8 +395,6 @@ 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})
@@ -405,23 +420,16 @@ 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:
-        if role == "Citizen":
-            completion = RAG(query, chat_history, role=role) 
-        else:
-            completion = RAG(query, chat_history, role=role)
+        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:
@@ -436,8 +444,6 @@ 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
@@ -464,7 +470,6 @@ 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"
@@ -475,7 +480,6 @@ 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"
@@ -490,14 +494,13 @@ 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 # Keep local sync
+                        question_count = user_to_update.question_count
                         await db.commit()
                     break
                     
@@ -512,9 +515,7 @@ 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)