Merge pull request #697 from CREDO23/implement-surfsense-docs-mentions

[Feat] Capture Google profile data, add user profile settings & document mentions picker improvements

surfsense_backend/alembic/versions/62_add_user_profile_columns.py

@@ -0,0 +1,72 @@
+"""Add display_name and avatar_url columns to user table
+
+This migration adds:
+- display_name column for user's full name from OAuth
+- avatar_url column for user's profile picture URL from OAuth
+
+Revision ID: 62
+Revises: 61
+"""
+
+from collections.abc import Sequence
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "62"
+down_revision: str | None = "61"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Add display_name and avatar_url columns to user table."""
+
+    # Add display_name column (nullable for existing users)
+    op.execute(
+        """
+        DO $$
+        BEGIN
+            IF NOT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'user' AND column_name = 'display_name'
+            ) THEN
+                ALTER TABLE "user" 
+                ADD COLUMN display_name VARCHAR;
+            END IF;
+        END$$;
+        """
+    )
+
+    # Add avatar_url column (nullable for existing users)
+    op.execute(
+        """
+        DO $$
+        BEGIN
+            IF NOT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'user' AND column_name = 'avatar_url'
+            ) THEN
+                ALTER TABLE "user" 
+                ADD COLUMN avatar_url VARCHAR;
+            END IF;
+        END$$;
+        """
+    )
+
+
+def downgrade() -> None:
+    """Remove display_name and avatar_url columns from user table."""
+
+    op.execute(
+        """
+        ALTER TABLE "user" 
+        DROP COLUMN IF EXISTS avatar_url;
+        """
+    )
+    op.execute(
+        """
+        ALTER TABLE "user" 
+        DROP COLUMN IF EXISTS display_name;
+        """
+    )

surfsense_backend/alembic/versions/63_add_message_author_id.py

@@ -0,0 +1,47 @@
+"""Add author_id column to new_chat_messages table
+
+Revision ID: 63
+Revises: 62
+"""
+
+from collections.abc import Sequence
+
+from alembic import op
+
+revision: str = "63"
+down_revision: str | None = "62"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Add author_id column to new_chat_messages table."""
+    op.execute(
+        """
+        DO $$
+        BEGIN
+            IF NOT EXISTS (
+                SELECT 1 FROM information_schema.columns
+                WHERE table_name = 'new_chat_messages' AND column_name = 'author_id'
+            ) THEN
+                ALTER TABLE new_chat_messages
+                ADD COLUMN author_id UUID REFERENCES "user"(id) ON DELETE SET NULL;
+                
+                CREATE INDEX IF NOT EXISTS ix_new_chat_messages_author_id
+                ON new_chat_messages(author_id);
+            END IF;
+        END$$;
+        """
+    )
+
+
+def downgrade() -> None:
+    """Remove author_id column from new_chat_messages table."""
+    op.execute(
+        """
+        DROP INDEX IF EXISTS ix_new_chat_messages_author_id;
+        ALTER TABLE new_chat_messages
+        DROP COLUMN IF EXISTS author_id;
+        """
+    )
+

surfsense_backend/app/db.py

@@ -413,8 +413,17 @@ class NewChatMessage(BaseModel, TimestampMixin):
         index=True,
     )
 
-    # Relationship
+    # Track who sent this message (for shared chats)
+    author_id = Column(
+        UUID(as_uuid=True),
+        ForeignKey("user.id", ondelete="SET NULL"),
+        nullable=True,
+        index=True,
+    )
+
+    # Relationships
     thread = relationship("NewChatThread", back_populates="messages")
+    author = relationship("User")
 
 
 class Document(BaseModel, TimestampMixin):
@@ -876,6 +885,10 @@ if config.AUTH_TYPE == "GOOGLE":
         )
         pages_used = Column(Integer, nullable=False, default=0, server_default="0")
 
+        # User profile from OAuth
+        display_name = Column(String, nullable=True)
+        avatar_url = Column(String, nullable=True)
+
 else:
 
     class User(SQLAlchemyBaseUserTableUUID, Base):
@@ -909,6 +922,10 @@ else:
         )
         pages_used = Column(Integer, nullable=False, default=0, server_default="0")
 
+        # User profile (can be set manually for non-OAuth users)
+        display_name = Column(String, nullable=True)
+        avatar_url = Column(String, nullable=True)
+
 
 engine = create_async_engine(DATABASE_URL)
 async_session_maker = async_sessionmaker(engine, expire_on_commit=False)

surfsense_backend/app/routes/new_chat_routes.py

@@ -411,11 +411,9 @@ async def get_thread_messages(
     Requires CHATS_READ permission.
     """
     try:
-        # Get thread with messages
+        # Get thread first
         result = await session.execute(
-            select(NewChatThread)
-            .options(selectinload(NewChatThread.messages))
-            .filter(NewChatThread.id == thread_id)
+            select(NewChatThread).filter(NewChatThread.id == thread_id)
         )
         thread = result.scalars().first()
 
@@ -434,6 +432,15 @@ async def get_thread_messages(
         # Check thread-level access based on visibility
         await check_thread_access(session, thread, user)
 
+        # Get messages with their authors loaded
+        messages_result = await session.execute(
+            select(NewChatMessage)
+            .options(selectinload(NewChatMessage.author))
+            .filter(NewChatMessage.thread_id == thread_id)
+            .order_by(NewChatMessage.created_at)
+        )
+        db_messages = messages_result.scalars().all()
+
         # Return messages in the format expected by assistant-ui
         messages = [
             NewChatMessageRead(
@@ -442,8 +449,11 @@ async def get_thread_messages(
                 role=msg.role,
                 content=msg.content,
                 created_at=msg.created_at,
+                author_id=msg.author_id,
+                author_display_name=msg.author.display_name if msg.author else None,
+                author_avatar_url=msg.author.avatar_url if msg.author else None,
             )
-            for msg in thread.messages
+            for msg in db_messages
         ]
 
         return ThreadHistoryLoadResponse(messages=messages)
@@ -782,6 +792,7 @@ async def append_message(
             thread_id=thread_id,
             role=message_role,
             content=message.content,
+            author_id=user.id,
         )
         session.add(db_message)
 

surfsense_backend/app/schemas/new_chat.py

@@ -38,6 +38,9 @@ class NewChatMessageRead(NewChatMessageBase, IDModel, TimestampModel):
     """Schema for reading a message."""
 
     thread_id: int
+    author_id: UUID | None = None
+    author_display_name: str | None = None
+    author_avatar_url: str | None = None
     model_config = ConfigDict(from_attributes=True)
 
 

surfsense_backend/app/schemas/users.py

@@ -6,6 +6,8 @@ from fastapi_users import schemas
 class UserRead(schemas.BaseUser[uuid.UUID]):
     pages_limit: int
     pages_used: int
+    display_name: str | None = None
+    avatar_url: str | None = None
 
 
 class UserCreate(schemas.BaseUserCreate):
@@ -13,4 +15,5 @@ class UserCreate(schemas.BaseUserCreate):
 
 
 class UserUpdate(schemas.BaseUserUpdate):
-    pass
+    display_name: str | None = None
+    avatar_url: str | None = None

surfsense_backend/app/users.py

@@ -1,6 +1,7 @@
 import logging
 import uuid
 
+import httpx
 from fastapi import Depends, Request, Response
 from fastapi.responses import JSONResponse, RedirectResponse
 from fastapi_users import BaseUserManager, FastAPIUsers, UUIDIDMixin, models
@@ -46,6 +47,71 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
     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: int | None = None,
+        refresh_token: str | None = None,
+        request: Request | None = None,
+        *,
+        associate_by_email: bool = False,
+        is_verified_by_default: bool = False,
+    ) -> User:
+        """
+        Override OAuth callback to capture Google profile data (name, avatar).
+        """
+        # Call parent implementation to create/get user
+        user = await super().oauth_callback(
+            oauth_name,
+            access_token,
+            account_id,
+            account_email,
+            expires_at,
+            refresh_token,
+            request,
+            associate_by_email=associate_by_email,
+            is_verified_by_default=is_verified_by_default,
+        )
+
+        # Fetch and store Google profile data if not already set
+        if oauth_name == "google" and (not user.display_name or not user.avatar_url):
+            try:
+                async with httpx.AsyncClient() as client:
+                    response = await client.get(
+                        "https://people.googleapis.com/v1/people/me",
+                        params={"personFields": "names,photos"},
+                        headers={"Authorization": f"Bearer {access_token}"},
+                    )
+                    response.raise_for_status()
+                    profile = response.json()
+
+                update_dict = {}
+
+                # Extract name from names array
+                names = profile.get("names", [])
+                if not user.display_name and names:
+                    display_name = names[0].get("displayName")
+                    if display_name:
+                        update_dict["display_name"] = display_name
+
+                # Extract photo URL from photos array
+                photos = profile.get("photos", [])
+                if not user.avatar_url and photos:
+                    photo_url = photos[0].get("url")
+                    if photo_url:
+                        update_dict["avatar_url"] = photo_url
+
+                if update_dict:
+                    user = await self.user_db.update(user, update_dict)
+
+            except Exception as e:
+                logger.warning(f"Failed to fetch Google profile: {e}")
+
+        return user
+
     async def on_after_register(self, user: User, request: Request | None = None):
         """
         Called after a user registers. Creates a default search space for the user