Update Model

api/auth.py

@@ -1,10 +1,15 @@
+# api/auth.py
 from fastapi_users import FastAPIUsers
 from fastapi_users.authentication import CookieTransport, JWTStrategy, AuthenticationBackend
 from fastapi_users.db import SQLAlchemyUserDatabase
 from httpx_oauth.clients.google import GoogleOAuth2
 from httpx_oauth.clients.github import GitHubOAuth2
-from api.database import SessionLocal
-from api.models import User, OAuthAccount
+from api.database import User, OAuthAccount, get_user_db
+from fastapi_users.manager import BaseUserManager, IntegerIDMixin
+from fastapi import Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+from fastapi_users.models import UP
+from typing import Optional
 import os
 import logging
 
@@ -46,9 +51,87 @@ if not all([GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, GITHUB_CLIENT_ID, GITHUB_CLI
 google_oauth_client = GoogleOAuth2(GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET)
 github_oauth_client = GitHubOAuth2(GITHUB_CLIENT_ID, GITHUB_CLIENT_SECRET)
 
+class UserManager(IntegerIDMixin, BaseUserManager[User, int]):
+    reset_password_token_secret = SECRET
+    verification_token_secret = SECRET
+
+    async def oauth_callback(
+        self,
+        oauth_name: str,
+        access_token: str,
+        account_id: str,
+        account_email: str,
+        expires_at: Optional[int] = None,
+        refresh_token: Optional[str] = None,
+        request: Optional[Request] = None,
+        *,
+        associate_by_email: bool = False,
+        is_verified_by_default: bool = False,
+    ) -> UP:
+        oauth_account_dict = {
+            "oauth_name": oauth_name,
+            "access_token": access_token,
+            "account_id": account_id,
+            "account_email": account_email,
+            "expires_at": expires_at,
+            "refresh_token": refresh_token,
+        }
+        oauth_account = OAuthAccount(**oauth_account_dict)
+        existing_oauth_account = await self.user_db.get_by_oauth_account(oauth_name, account_id)
+        if existing_oauth_account is not None:
+            return await self.on_after_login(existing_oauth_account.user, request)
+
+        if associate_by_email:
+            user = await self.user_db.get_by_email(account_email)
+            if user is not None:
+                oauth_account.user_id = user.id
+                await self.user_db.add_oauth_account(oauth_account)
+                return await self.on_after_login(user, request)
+
+        user_dict = {
+            "email": account_email,
+            "hashed_password": self.password_helper.hash("dummy_password"),
+            "is_active": True,
+            "is_verified": is_verified_by_default,
+        }
+        user = await self.user_db.create(user_dict)
+        oauth_account.user_id = user.id
+        await self.user_db.add_oauth_account(oauth_account)
+        return await self.on_after_login(user, request)
+
+async def get_user_manager(user_db: SQLAlchemyUserDatabase = Depends(get_user_db)):
+    yield UserManager(user_db)
+
+from fastapi_users.router.oauth import get_oauth_router
+
+google_oauth_router = get_oauth_router(
+    google_oauth_client,
+    auth_backend,
+    get_user_manager,
+    associate_by_email=True,
+    redirect_url="https://mgzon-mgzon-app.hf.space/auth/google/callback",
+)
+
+github_oauth_router = get_oauth_router(
+    github_oauth_client,
+    auth_backend,
+    get_user_manager,
+    associate_by_email=True,
+    redirect_url="https://mgzon-mgzon-app.hf.space/auth/github/callback",
+)
+
 fastapi_users = FastAPIUsers[User, int](
-    lambda: SQLAlchemyUserDatabase(User, SessionLocal(), oauth_account_table=OAuthAccount),
+    get_user_db,
     [auth_backend],
 )
 
 current_active_user = fastapi_users.current_user(active=True, optional=True)
+
+def get_auth_router(app: FastAPI):
+    app.include_router(google_oauth_router, prefix="/auth/google", tags=["auth"])
+    app.include_router(github_oauth_router, prefix="/auth/github", tags=["auth"])
+    app.include_router(fastapi_users.get_auth_router(auth_backend), prefix="/auth/jwt", tags=["auth"])
+    app.include_router(fastapi_users.get_register_router(UserRead, UserCreate), prefix="/auth", tags=["auth"])
+    app.include_router(fastapi_users.get_reset_password_router(), prefix="/auth", tags=["auth"])
+    app.include_router(fastapi_users.get_verify_router(UserRead), prefix="/auth", tags=["auth"])
+    app.include_router(fastapi_users.get_users_router(UserRead, UserUpdate), prefix="/users", tags=["users"])

api/database.py

@@ -1,13 +1,20 @@
 # api/database.py
 import os
+from sqlalchemy import Column, String, Integer, ForeignKey, DateTime, Boolean
 from sqlalchemy import create_engine
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import sessionmaker, relationship
+from sqlalchemy.sql import func
+from fastapi_users.db import SQLAlchemyBaseUserTable, SQLAlchemyUserDatabase
+from sqlalchemy.orm import Session
+from typing import AsyncGenerator
+from fastapi import Depends
+from datetime import datetime
 
 # جلب URL قاعدة البيانات من المتغيرات البيئية
 SQLALCHEMY_DATABASE_URL = os.getenv("SQLALCHEMY_DATABASE_URL")
 if not SQLALCHEMY_DATABASE_URL:
-    raise ValueError("SQLALCHEMY_DATABASE_URL is not set in environment variables.")
+    raise ValueValue("SQLALCHEMY_DATABASE_URL is not set in environment variables.")
 
 # إنشاء المحرك
 engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
@@ -18,10 +25,64 @@ SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 # قاعدة أساسية للنماذج
 Base = declarative_base()
 
-# دالة للحصول على الجلسة
+class OAuthAccount(Base):
+    __tablename__ = "oauth_account"
+
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(Integer, ForeignKey("user.id"), nullable=False)
+    oauth_name = Column(String, nullable=False)
+    access_token = Column(String, nullable=False)
+    expires_at = Column(Integer, nullable=True)
+    refresh_token = Column(String, nullable=True)
+    account_id = Column(String, index=True, nullable=False)
+    account_email = Column(String, nullable=False)
+
+    user = relationship("User", back_populates="oauth_accounts")
+
+class User(SQLAlchemyBaseUserTable[int], Base):
+    __tablename__ = "user"
+
+    id = Column(Integer, primary_key=True, index=True)
+    email = Column(String, unique=True, index=True, nullable=False)
+    hashed_password = Column(String, nullable=False)
+    is_active = Column(Boolean, default=True)
+    is_superuser = Column(Boolean, default=False)
+    is_verified = Column(Boolean, default=False)
+    display_name = Column(String, nullable=True)
+    preferred_model = Column(String, nullable=True)
+    job_title = Column(String, nullable=True)
+    education = Column(String, nullable=True)
+    interests = Column(String, nullable=True)
+    additional_info = Column(String, nullable=True)
+    conversation_style = Column(String, nullable=True)
+    oauth_accounts = relationship("OAuthAccount", back_populates="user", cascade="all, delete-orphan")
+
+class Conversation(Base):
+    __tablename__ = "conversation"
+
+    id = Column(Integer, primary_key=True, index=True)
+    conversation_id = Column(String, unique=True, index=True)
+    user_id = Column(Integer, ForeignKey("user.id"))
+    title = Column(String)
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+    updated_at = Column(DateTime(timezone=True), onupdate=func.now())
+
+class Message(Base):
+    __tablename__ = "message"
+
+    id = Column(Integer, primary_key=True, index=True)
+    conversation_id = Column(Integer, ForeignKey("conversation.id"))
+    role = Column(String)
+    content = Column(String)
+
 def get_db():
     db = SessionLocal()
     try:
         yield db
     finally:
         db.close()
+
+Base.metadata.create_all(bind=engine)
+
+async def get_user_db(session: Session = Depends(get_db)):
+    yield SQLAlchemyUserDatabase(session, User, OAuthAccount)

api/models.py

@@ -1,4 +1,5 @@
-from fastapi_users.db import SQLAlchemyBaseUserTable, SQLAlchemyBaseOAuthAccountTable
+# models.py
+from fastapi_users.db import SQLAlchemyBaseUserTable
 from sqlalchemy import Column, Integer, String, Boolean, Text, ForeignKey, DateTime
 from sqlalchemy.orm import relationship
 from sqlalchemy.ext.declarative import declarative_base
@@ -11,9 +12,16 @@ import uuid
 Base = declarative_base()
 
 # جدول OAuth Accounts لتخزين بيانات تسجيل الدخول الخارجي
-class OAuthAccount(SQLAlchemyBaseOAuthAccountTable, Base):
+class OAuthAccount(Base):
     __tablename__ = "oauth_accounts"
     id = Column(Integer, primary_key=True)
+    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
+    oauth_name = Column(String, nullable=False)
+    access_token = Column(String, nullable=False)
+    expires_at = Column(Integer, nullable=True)
+    refresh_token = Column(String, nullable=True)
+    account_id = Column(String, index=True, nullable=False)
+    account_email = Column(String, nullable=False)
     user = relationship("User", back_populates="oauth_accounts")
 
 # نموذج المستخدم

main.py

@@ -1,3 +1,4 @@
+# main.py
 # SPDX-FileCopyrightText: Hadad <[email protected]>
 # SPDX-License-Identifier: Apache-2.0
 
@@ -12,9 +13,8 @@ from starlette.middleware.sessions import SessionMiddleware
 from fastapi.openapi.docs import get_swagger_ui_html
 from fastapi.middleware.cors import CORSMiddleware
 from api.endpoints import router as api_router
-from api.auth import fastapi_users, auth_backend, current_active_user, google_oauth_client, github_oauth_client
+from api.auth import fastapi_users, auth_backend, current_active_user, get_auth_router
 from api.database import get_db, engine, Base
-from api.models import User, UserRead, UserCreate, Conversation
 from motor.motor_asyncio import AsyncIOMotorClient
 from pydantic import BaseModel
 from typing import List
@@ -25,8 +25,8 @@ from sqlalchemy.orm import Session
 from pathlib import Path
 from hashlib import md5
 from datetime import datetime
-import re
 from httpx_oauth.exceptions import GetIdEmailError
+from api.models import UserRead, UserCreate, Conversation, UserUpdate
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -79,6 +79,7 @@ CONCURRENCY_LIMIT = int(os.getenv("CONCURRENCY_LIMIT", 20))
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     await setup_mongo_index()
+    Base.metadata.create_all(bind=engine)  # Create tables on startup
     yield
 
 app = FastAPI(title="MGZon Chatbot API", lifespan=lifespan)
@@ -86,9 +87,6 @@ app = FastAPI(title="MGZon Chatbot API", lifespan=lifespan)
 # Add SessionMiddleware
 app.add_middleware(SessionMiddleware, secret_key=JWT_SECRET)
 
-# Create SQLAlchemy tables
-Base.metadata.create_all(bind=engine)
-
 # Mount static files
 os.makedirs("static", exist_ok=True)
 app.mount("/static", StaticFiles(directory="static"), name="static")
@@ -106,42 +104,8 @@ app.add_middleware(
 )
 
 # Include routers
-app.include_router(
-    fastapi_users.get_auth_router(auth_backend),
-    prefix="/auth/jwt",
-    tags=["auth"],
-)
-app.include_router(
-    fastapi_users.get_register_router(UserRead, UserCreate),
-    prefix="/auth",
-    tags=["auth"],
-)
-app.include_router(
-    fastapi_users.get_users_router(UserRead, UserCreate),
-    prefix="/users",
-    tags=["users"],
-)
-app.include_router(
-    fastapi_users.get_oauth_router(
-        google_oauth_client,
-        auth_backend,
-        JWT_SECRET,
-        redirect_url="https://mgzon-mgzon-app.hf.space/auth/google/callback"
-    ),
-    prefix="/auth/google",
-    tags=["auth"],
-)
-app.include_router(
-    fastapi_users.get_oauth_router(
-        github_oauth_client,
-        auth_backend,
-        JWT_SECRET,
-        redirect_url="https://mgzon-mgzon-app.hf.space/auth/github/callback"
-    ),
-    prefix="/auth/github",
-    tags=["auth"],
-)
 app.include_router(api_router)
+get_auth_router(app)  # Add OAuth and auth routers
 
 # Debug routes endpoint
 @app.get("/debug/routes", response_class=PlainTextResponse)
@@ -153,81 +117,6 @@ async def debug_routes():
         routes.append(f"{methods} {path}")
     return "\n".join(sorted(routes))
 
-# OAuth callbacks
-@app.get("/auth/google/callback", response_class=RedirectResponse)
-async def google_oauth_callback(
-    request: Request,
-    code: str = Query(...),
-    state: str = Query(...),
-    db: Session = Depends(get_db)
-):
-    try:
-        logger.info("Processing Google OAuth callback")
-        # Exchange code for access token
-        token_data = await google_oauth_client.get_access_token(code, "https://mgzon-mgzon-app.hf.space/auth/google/callback")
-        logger.info(f"Google OAuth token received: {token_data}")
-        # Get user info
-        user_info = await google_oauth_client.get_id_email(token_data["access_token"])
-        logger.info(f"Google user info: {user_info}")
-        # Create or update user
-        user = await fastapi_users.oauth_callback(
-            oauth_name="google",
-            access_token=token_data["access_token"],
-            account_id=user_info["id"],
-            account_email=user_info["email"],
-            expires_at=token_data.get("expires_at"),
-            refresh_token=token_data.get("refresh_token"),
-            request=request,
-            db=db
-        )
-        logger.info("Google OAuth user processed, creating session")
-        # Create JWT token
-        token = await auth_backend.get_login_response(user, request)
-        logger.info("Google OAuth callback processed, redirecting to /chat")
-        response = RedirectResponse(url="/chat", status_code=302)
-        response.set_cookie("Authorization", f"Bearer {token.access_token}", httponly=True)
-        return response
-    except Exception as e:
-        logger.error(f"Google OAuth callback error: {str(e)}")
-        return RedirectResponse(url=f"/login?error=Google%20OAuth%20failed:%20{str(e)}", status_code=302)
-
-@app.get("/auth/github/callback", response_class=RedirectResponse)
-async def github_oauth_callback(
-    request: Request,
-    code: str = Query(...),
-    state: str = Query(...),
-    db: Session = Depends(get_db)
-):
-    try:
-        logger.info("Processing GitHub OAuth callback")
-        # Exchange code for access token
-        token_data = await github_oauth_client.get_access_token(code, "https://mgzon-mgzon-app.hf.space/auth/github/callback")
-        logger.info(f"GitHub OAuth token received: {token_data}")
-        # Get user info
-        user_info = await github_oauth_client.get_id_email(token_data["access_token"])
-        logger.info(f"GitHub user info: {user_info}")
-        # Create or update user
-        user = await fastapi_users.oauth_callback(
-            oauth_name="github",
-            access_token=token_data["access_token"],
-            account_id=user_info["id"],
-            account_email=user_info["email"],
-            expires_at=token_data.get("expires_at"),
-            refresh_token=token_data.get("refresh_token"),
-            request=request,
-            db=db
-        )
-        logger.info("GitHub OAuth user processed, creating session")
-        # Create JWT token
-        token = await auth_backend.get_login_response(user, request)
-        logger.info("GitHub OAuth callback processed, redirecting to /chat")
-        response = RedirectResponse(url="/chat", status_code=302)
-        response.set_cookie("Authorization", f"Bearer {token.access_token}", httponly=True)
-        return response
-    except Exception as e:
-        logger.error(f"GitHub OAuth callback error: {str(e)}")
-        return RedirectResponse(url=f"/login?error=GitHub%20OAuth%20failed:%20{str(e)}", status_code=302)
-
 # Custom middleware for 404 and 500 errors
 class NotFoundMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):