File size: 5,875 Bytes
fae4e5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
"""
HuggingFace Authentication for MockTraceMind
Using Gradio's built-in OAuth support (simpler than manual OAuth)
"""

import os
import gradio as gr
from typing import Optional
from dataclasses import dataclass

# Development mode flag - set DISABLE_OAUTH=true to skip OAuth for local dev
DEV_MODE = os.getenv("DISABLE_OAUTH", "false").lower() in ("true", "1", "yes")


@dataclass
class User:
    """Authenticated user information"""
    username: str
    name: str
    avatar_url: str
    token: str

    @classmethod
    def from_oauth(cls, token: gr.OAuthToken, profile: gr.OAuthProfile) -> "User":
        """Create User from Gradio OAuth objects"""
        return cls(
            username=profile.username,
            name=profile.name,
            avatar_url=profile.picture,
            token=str(token)
        )

    @classmethod
    def create_dev_user(cls) -> "User":
        """Create a mock user for development mode"""
        return cls(
            username=os.getenv("DEV_USERNAME", "dev_user"),
            name=os.getenv("DEV_NAME", "Development User"),
            avatar_url="https://huggingface.co/avatars/default-avatar.png",
            token="dev_token_12345"
        )


def is_authenticated(token: gr.OAuthToken | None, profile: gr.OAuthProfile | None) -> bool:
    """
    Check if user is authenticated

    Args:
        token: OAuth token from Gradio
        profile: OAuth profile from Gradio

    Returns:
        True if both token and profile are valid, or if in dev mode
    """
    # In dev mode, always consider authenticated
    if DEV_MODE:
        return True

    return token is not None and profile is not None


def get_user_info(token: gr.OAuthToken | None, profile: gr.OAuthProfile | None) -> Optional[User]:
    """
    Get user information from OAuth objects

    Args:
        token: OAuth token from Gradio
        profile: OAuth profile from Gradio

    Returns:
        User object if authenticated, None otherwise
    """
    if not is_authenticated(token, profile):
        return None

    # In dev mode, return mock user
    if DEV_MODE:
        return User.create_dev_user()

    return User.from_oauth(token, profile)


def create_login_handler(on_login_success=None, on_login_failure=None):
    """
    Create a login handler function for Gradio LoginButton

    Args:
        on_login_success: Callback function called when login succeeds
        on_login_failure: Callback function called when login fails

    Returns:
        Handler function compatible with Gradio LoginButton.click()
    """
    def handle_login(token: gr.OAuthToken | None, profile: gr.OAuthProfile | None):
        if is_authenticated(token, profile):
            user = get_user_info(token, profile)
            if on_login_success:
                return on_login_success(user)
            return user
        else:
            if on_login_failure:
                return on_login_failure()
            return None

    return handle_login


def require_auth(func):
    """
    Decorator to require authentication for a function

    Usage:
        @require_auth
        def my_function(user: User, other_args...):
            # user is guaranteed to be valid User object
            pass
    """
    def wrapper(token: gr.OAuthToken | None, profile: gr.OAuthProfile | None, *args, **kwargs):
        if not is_authenticated(token, profile):
            gr.Warning("Please log in to Hugging Face to access this feature!")
            return None

        user = get_user_info(token, profile)
        return func(user, *args, **kwargs)

    return wrapper


# UI component helpers
def create_login_button(visible: bool = True) -> gr.LoginButton:
    """
    Create a styled HuggingFace login button
    Automatically hidden in dev mode
    """
    # Hide login button in dev mode
    if DEV_MODE:
        visible = False

    return gr.LoginButton(visible=visible)


def create_user_info_display(user: Optional[User]) -> str:
    """
    Create HTML for user info display

    Args:
        user: User object or None

    Returns:
        HTML string for display
    """
    if user is None:
        # In dev mode, don't show login prompt
        if DEV_MODE:
            return """
            <div style="text-align: center; padding: 10px; border: 2px solid #ffa500; border-radius: 10px; background-color: #fff4e6;">
                <strong>🛠️ Development Mode</strong>
                <p style="margin: 5px 0 0 0; font-size: 0.9em;">OAuth disabled for local testing</p>
            </div>
            """

        return """
        <div style="text-align: center; padding: 20px; border: 2px dashed #ccc; border-radius: 10px;">
            <h3>🔒 Login Required</h3>
            <p>Please log in with your Hugging Face account to access TraceMind</p>
        </div>
        """

    # Add dev mode badge if in dev mode
    dev_badge = ""
    if DEV_MODE:
        dev_badge = '<span style="background: #ffa500; color: white; padding: 2px 8px; border-radius: 4px; font-size: 0.8em; margin-left: 10px;">DEV</span>'

    return f"""
    <div style="display: flex; align-items: center; padding: 10px; border: 1px solid #e0e0e0; border-radius: 8px;">
        <img src="{user.avatar_url}" alt="{user.name}"
             style="width: 48px; height: 48px; border-radius: 50%; margin-right: 15px;">
        <div>
            <strong>{user.name}</strong>{dev_badge}<br>
            <small style="color: #666;">@{user.username}</small>
        </div>
    </div>
    """


def create_auth_warning(message: str = "Please login first") -> str:
    """Create a warning message for unauthenticated users"""
    return f"""
    <div style="text-align: center; padding: 20px; border: 2px solid #ff6b6b; border-radius: 10px; background-color: #ffe0e0;">
        <h3>⚠️ Authentication Required</h3>
        <p>{message}</p>
    </div>
    """