merging code with dev

Author: AnkitTukatek

backend/alembic/versions/7d819bf06aab_add_chat_feedback_timestamp.py

@@ -0,0 +1,33 @@
+"""add_chat_feedback_timestamp
+
+Revision ID: 7d819bf06aab
+Revises: 0ba3b4464c70
+Create Date: 2025-08-14 03:36:05.961388
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "7d819bf06aab"
+down_revision = "0ba3b4464c70"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "chat_feedback",
+        sa.Column(
+            "feedback_timestamp",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=True,
+        ),
+    )
+
+def downgrade() -> None:
+    op.drop_column("chat_feedback", "feedback_timestamp")

backend/onyx/auth/unified_oauth_callback.py

@@ -322,7 +322,7 @@
 
 #     return router
 
-
+from onyx.db.engine.async_sql_engine import get_async_session 
 import httpx
 from fastapi import APIRouter, Depends, Request, HTTPException, status
 from fastapi.responses import RedirectResponse
@@ -353,12 +353,17 @@
     is_domain_allowed_for_oauth,
     get_allowed_oauth_domains
 )
+from onyx.db.models import User
+from onyx.auth.sso_data_db import get_sso_configurations_from_db
 import jwt
 import json
+from sqlalchemy import select
 
 STATE_TOKEN_AUDIENCE = "fastapi-users:oauth-state"
 logger = setup_logger()
 
+GRAPH_ROOT = "https://graph.microsoft.com/v1.0"
+
 async def get_google_user_info(access_token: str) -> dict:
     async with httpx.AsyncClient() as client:
         response = await client.get(
@@ -367,17 +372,83 @@ async def get_google_user_info(access_token: str) -> dict:
         )
         response.raise_for_status()
         return response.json()
-    
-async def get_microsoft_user_info(access_token: str) -> dict:
-    logger.info(f"accesssssssssssssssssssss{access_token}")
+async def _make_graph_request(access_token: str, endpoint: str) -> dict:
+    """Helper function to make Microsoft Graph API requests"""
+    logger.info(f"Making request to {endpoint} with token: {access_token}")
+
     async with httpx.AsyncClient() as client:
         response = await client.get(
-            "https://graph.microsoft.com/v1.0/me",
+            f"{GRAPH_ROOT}/{endpoint}",
             headers={"Authorization": f"Bearer {access_token}"}
         )
         response.raise_for_status()
-        logger.info(f"response:{response.json()}")
-        return response.json()
+        result = response.json()
+        logger.info(f"Response: {result}")
+        return result
+
+
+async def get_microsoft_user_info(access_token: str) -> dict:
+    return await _make_graph_request(access_token, "me")
+
+
+async def check_microsoft_account(access_token: str) -> dict:
+    return await _make_graph_request(access_token, "me/appRoleAssignments")
+
+
+async def check_entra_account_details(access_token: str) -> dict:
+    from onyx.db.engine.async_sql_engine import get_async_session    
+    db_generator = get_async_session()
+    db = await anext(db_generator)
+    client_id = None
+    
+    sso_configs = await get_sso_configurations_from_db(db)
+    for provider, creds in sso_configs.items():
+        if creds.get("client_id") and creds.get("client_secret"):
+            client_id = creds["client_id"]
+            break  # stop after finding the first match
+    user_account_details = await _make_graph_request(access_token, "me/appRoleAssignments")
+
+    account_details = await _make_graph_request(
+        access_token,
+        f"servicePrincipals?$filter=appId eq '{client_id}'"
+    )
+    user_roles = user_account_details.get("value", [])
+    accounts = account_details.get("value", [])
+    logger.info(f"User account details 3 : f{user_roles} f{accounts}")
+    if not accounts:
+        return {"account_exists": False, "account_role_check": False}
+
+    first_account = accounts[0]
+
+    # Check if any user role matches the account resourceId
+    account_exists = [
+        role for role in user_roles if role.get("resourceId") == first_account.get("id")
+    ]
+
+    account_role_check = account_exists  # Same filter as above
+    logger.info(f"User account details 3 : f{account_role_check} ")
+    if account_exists and account_role_check:
+        # Fetch role type (appRole value)
+        app_roles = first_account.get("appRoles", [])
+        role_type = "User"
+        if app_roles and account_role_check:
+            for role in app_roles:
+                if role.get("id") == account_role_check[0].get("appRoleId"):
+                    role_type = role.get("value")
+                    break
+        logger.info(f"User account details6: {account_exists}")
+        return {
+            "account_exists": True,
+            "account_role_check": True,
+            "account_details": first_account,
+            "role_type": role_type
+        }
+
+    return {
+        "account_exists": False,
+        "account_role_check": False
+    }
+
 
 def get_unified_oauth_callback_router(
     auth_backend: AuthenticationBackend,
@@ -594,11 +665,21 @@ async def callback(
                     account_email = user_info.get("email")
                     account_id = user_info.get("sub")
                 else:  # microsoft
-                    user_info = await get_microsoft_user_info(token["access_token"])
+                    user_info = await get_microsoft_user_info(token["access_token"])            
+                    logger.info(f"Fetching Microsoft user info... {token['access_token']}")
+
+                    entra_account_details = await check_entra_account_details(token["access_token"])
+                    logger.info(f"User account details: {entra_account_details}")
+                    # Validate account existence
+                    if not entra_account_details or not entra_account_details.get("account_exists", False):
+                        raise Exception("Account does not exist in Microsoft Entra")
+
+                    # Extract account details
                     account_email = user_info.get("userPrincipalName") or user_info.get("mail")
                     account_id = user_info.get("id")
+
 
-                logger.info(f"Retrieved user info - email: {account_email}, id: {account_id}")
+                    logger.info(f"Retrieved user info - email: {account_email}, id: {account_id}")
 
             except Exception as e:
                 logger.error(f"Failed to fetch user info: {str(e)}")
@@ -618,7 +699,16 @@ async def callback(
             # Check if this OAuth registration should be allowed
             oauth_allowed = True if provider == "microsoft" else should_allow_oauth_registration(account_email, provider)
 
-            
+            role_Type_check1 = (
+                    "BASIC"
+                    if provider == "google"
+                    else "ADMIN" if entra_account_details.get("role_type") == "Admin" else "BASIC"
+                )
+            role_Type_check = (
+                    "basic"
+                    if provider == "google"
+                    else "admin" if entra_account_details.get("role_type") == "Admin" else "basic"
+                )
             if not oauth_allowed:
                 logger.error(f"OAuth registration not allowed for {account_email} with provider {provider}")
 
@@ -653,7 +743,9 @@ async def callback(
             request.state.referral_source = referral_source
             request.state.oauth_provider = provider
             request.state.is_oauth_flow = True
-            request.state.user_role = "basic"  # Set default role for OAuth users
+            # request.state.user_role = "admin" if entra_account_details.get("role_type") == "Admin" else "basic"
+            request.state.user_role = role_Type_check
+            # Set default role for OAuth users
 
             # Perform OAuth callback to create/login user
             try:
@@ -670,9 +762,17 @@ async def callback(
                     request=request,
                     associate_by_email=True,
                     is_verified_by_default=True,
+                    # role_type=role_Type_check
                 )
-                logger.info(f"OAuth callback successful for user: {user.email}")
+                logger.info(f"OAuth callback successful for user: {user.email} ")
+                logger.info(f"OAuth user: {user} ")
+                role_update_success = await update_user_role_in_db(role_Type_check1, account_email)
 
+                if role_update_success:
+                    logger.info(f"Successfully updated user {account_email} role to {role_Type_check}")
+                else:
+                    logger.warning(f"Failed to update role for user {account_email}, but continuing with login")
+
             except UserAlreadyExists:
                 logger.error("User already exists")
                 raise HTTPException(status_code=400, detail="User already exists")
@@ -756,3 +856,39 @@ async def callback(
             raise HTTPException(status_code=500, detail="Internal server error during OAuth callback")
 
     return router
+
+async def update_user_role_in_db(role: str, email: str):
+
+    try:
+        db_generator = get_async_session()
+        db = await anext(db_generator)
+        
+        try:
+            result = await db.execute(
+                select(User).where(User.email == email)
+            )
+            logger.info(f"Role type {result}")
+            user = result.unique().scalar_one_or_none()
+            
+            if not user:
+                logger.error(f"User with ID {User.email} not found for role update")
+                return False
+            
+            old_role = user.role
+            user.role = role
+            
+            await db.commit()
+            
+            logger.info(f"Successfully updated user {email} role from {old_role} to {role}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to update user role for {email}: {str(e)}")
+            await db.rollback()
+            return False
+        finally:
+            await db.close()
+            
+    except Exception as e:
+        logger.error(f"Database session error when updating role for {email}: {str(e)}")
+        return False

backend/onyx/auth/users.py

@@ -401,7 +401,9 @@ async def oauth_callback(
         *,
         associate_by_email: bool = False,
         is_verified_by_default: bool = False,
+        role_type: str = UserRole.BASIC,
     ) -> User:
+        logger.info(f"role_typeeeeeeee {role_type}")
         referral_source = (
             getattr(request.state, "referral_source", None) if request else None
         )
@@ -442,6 +444,7 @@ async def oauth_callback(
                 "account_email": account_email,
                 "expires_at": expires_at,
                 "refresh_token": refresh_token,
+                # "role": role_type,
             }
 
             user: User | None = None
@@ -473,6 +476,7 @@ async def oauth_callback(
                         "email": account_email,
                         "hashed_password": self.password_helper.hash(password),
                         "is_verified": is_verified_by_default,
+                        # "role": role_type,
                     }
 
                     user = await self.user_db.create(user_dict)
@@ -518,7 +522,8 @@ async def oauth_callback(
                     user,
                     {
                         "is_verified": is_verified_by_default,
-                        "role": UserRole.BASIC,
+                        # "role": UserRole.BASIC,
+                        #  "role": role_type,
                     },
                 )
 

backend/onyx/db/chat.py

@@ -1030,12 +1030,34 @@ def get_retrieval_docs_from_search_docs(
     if sort_by_score:
         top_documents = sorted(top_documents, key=lambda doc: doc.score, reverse=True)  # type: ignore
     return RetrievalDocs(top_documents=top_documents)
+from onyx.db.models import ChatMessageFeedback
 
 
 def translate_db_message_to_chat_message_detail(
     chat_message: ChatMessage,
     remove_doc_content: bool = False,
+    db_session: Session = None,  # ADD THIS PARAMETER
 ) -> ChatMessageDetail:
+    
+    # ADD THIS: Get chat feedback for this message
+    chat_feedback = []
+    if db_session:
+        # Replace 'ChatFeedback' with your actual feedback model name
+        feedback_records = db_session.query(ChatMessageFeedback).filter(
+            ChatMessageFeedback.chat_message_id == chat_message.id
+        ).all()
+        
+        chat_feedback = [
+            {
+                "id": feedback.id,
+                "is_positive": feedback.is_positive,
+                "feedback_text": feedback.feedback_text,
+                "chat_message_id": feedback.chat_message_id,
+                "feedback_timestamp":feedback.feedback_timestamp
+            }
+            for feedback in feedback_records
+        ]
+    
     chat_msg_detail = ChatMessageDetail(
         chat_session_id=chat_message.chat_session_id,
         message_id=chat_message.id,
@@ -1067,6 +1089,7 @@ def translate_db_message_to_chat_message_detail(
         refined_answer_improvement=chat_message.refined_answer_improvement,
         is_agentic=chat_message.is_agentic,
         error=chat_message.error,
+        chat_feedback=chat_feedback,  # ADD THIS LINE
     )
 
     return chat_msg_detail

backend/onyx/db/models.py

@@ -5,7 +5,6 @@
 from typing import NotRequired
 from typing import Optional
 from uuid import uuid4
-
 from pydantic import BaseModel
 from sqlalchemy.orm import validates
 from typing_extensions import TypedDict  # noreorder
@@ -2336,6 +2335,11 @@ class ChatMessageFeedback(Base):
     required_followup: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
     feedback_text: Mapped[str | None] = mapped_column(Text, nullable=True)
     predefined_feedback: Mapped[str | None] = mapped_column(String, nullable=True)
+    feedback_timestamp: Mapped[datetime.datetime | None] = mapped_column(
+        DateTime(timezone=True), 
+        nullable=True, 
+        server_default=func.now()  # Set at database level
+    )
 
     chat_message: Mapped[ChatMessage] = relationship(
         "ChatMessage",