Spaces:
Running
Running
Debug
Browse files- app/main.py +129 -204
app/main.py
CHANGED
@@ -9,117 +9,87 @@ import logging
|
|
9 |
from contextlib import asynccontextmanager
|
10 |
|
11 |
from fastapi import FastAPI, Depends # Import FastAPI itself
|
12 |
-
# --- Import necessary items from database.py ---
|
13 |
from .database import connect_db, disconnect_db, database, metadata, users
|
14 |
from .api import router as api_router
|
15 |
from . import schemas, auth, dependencies
|
16 |
-
from .websocket import manager
|
17 |
-
|
18 |
from sqlalchemy.schema import CreateTable
|
19 |
-
from sqlalchemy.dialects import sqlite
|
20 |
|
21 |
-
# Configure logging
|
22 |
logging.basicConfig(level=logging.INFO)
|
23 |
logger = logging.getLogger(__name__)
|
24 |
|
25 |
-
|
26 |
-
API_BASE_URL = "http://127.0.0.1:7860/api" # Adjust if needed
|
27 |
|
28 |
-
# ---
|
29 |
@asynccontextmanager
|
30 |
async def lifespan(app: FastAPI):
|
31 |
-
#
|
32 |
logger.info("Application startup: Connecting DB...")
|
33 |
await connect_db()
|
34 |
logger.info("Application startup: DB Connected. Checking/Creating tables...")
|
35 |
-
if database.is_connected:
|
36 |
try:
|
37 |
-
# 1. Check if table exists using the async connection
|
38 |
-
# For SQLite, check the sqlite_master table
|
39 |
check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;"
|
40 |
table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name})
|
41 |
-
|
42 |
if not table_exists:
|
43 |
logger.info(f"Table '{users.name}' not found, attempting creation using async connection...")
|
44 |
-
|
45 |
-
# 2. Generate the CREATE TABLE statement using SQLAlchemy DDL compiler
|
46 |
-
# We need a dialect object for the compiler
|
47 |
-
# Infer dialect from the database URL if possible, default to SQLite for this example
|
48 |
-
dialect = sqlite.dialect() # Or determine dynamically based on database.url
|
49 |
create_table_stmt = str(CreateTable(users).compile(dialect=dialect))
|
50 |
-
logger.debug(f"Generated CREATE TABLE statement: {create_table_stmt}")
|
51 |
-
|
52 |
-
# 3. Execute the CREATE TABLE statement via the async database connection
|
53 |
await database.execute(query=create_table_stmt)
|
54 |
logger.info(f"Table '{users.name}' created successfully via async connection.")
|
55 |
-
|
56 |
-
# 4. Optional: Verify again immediately
|
57 |
table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name})
|
58 |
-
if table_exists_after:
|
59 |
-
|
60 |
-
else:
|
61 |
-
logger.error(f"Table '{users.name}' verification FAILED after creation attempt!")
|
62 |
-
|
63 |
else:
|
64 |
logger.info(f"Table '{users.name}' already exists (checked via async connection).")
|
65 |
-
|
66 |
except Exception as db_setup_err:
|
67 |
logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}")
|
68 |
-
# Consider whether to halt startup here depending on severity
|
69 |
else:
|
70 |
logger.error("CRITICAL: Database connection failed, skipping table setup.")
|
71 |
-
|
72 |
logger.info("Application startup: DB setup phase complete.")
|
73 |
yield
|
74 |
-
# Shutdown
|
75 |
logger.info("Application shutdown: Disconnecting DB...")
|
76 |
await disconnect_db()
|
77 |
logger.info("Application shutdown: DB Disconnected.")
|
78 |
|
79 |
-
# Create the main FastAPI app instance that Gradio will use
|
80 |
-
# We attach our API routes to this instance.
|
81 |
-
app = FastAPI(lifespan=lifespan)
|
82 |
-
app.include_router(api_router, prefix="/api") # Mount API routes under /api
|
83 |
-
|
84 |
-
# --- Gradio UI Definition ---
|
85 |
-
|
86 |
-
# Store websocket connection globally (or within a class) for the Gradio app instance
|
87 |
-
# This is tricky because Gradio re-runs functions. State management is key.
|
88 |
-
# We'll connect the WebSocket *after* login and store the connection task/info in gr.State.
|
89 |
|
90 |
-
# ---
|
|
|
|
|
91 |
|
|
|
92 |
async def make_api_request(method: str, endpoint: str, **kwargs):
|
93 |
async with httpx.AsyncClient() as client:
|
94 |
url = f"{API_BASE_URL}{endpoint}"
|
95 |
try:
|
96 |
response = await client.request(method, url, **kwargs)
|
97 |
-
response.raise_for_status()
|
98 |
return response.json()
|
99 |
except httpx.RequestError as e:
|
100 |
logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}")
|
101 |
return {"error": f"Network error contacting API: {e}"}
|
102 |
except httpx.HTTPStatusError as e:
|
103 |
logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
|
104 |
-
try:
|
105 |
-
|
106 |
-
except json.JSONDecodeError:
|
107 |
-
detail = e.response.text
|
108 |
return {"error": f"API Error: {detail}"}
|
109 |
except Exception as e:
|
110 |
logger.error(f"Unexpected error during API call: {e}")
|
111 |
return {"error": f"An unexpected error occurred: {str(e)}"}
|
112 |
|
113 |
# --- WebSocket handling within Gradio ---
|
114 |
-
|
115 |
-
async def listen_to_websockets(token: str,
|
116 |
-
"""Connects to WS and updates state list when a message arrives."""
|
117 |
ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
|
118 |
logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")
|
119 |
|
120 |
if not token:
|
121 |
logger.warning(f"[{ws_listener_id}] No token provided. Listener task exiting.")
|
122 |
-
|
|
|
123 |
|
124 |
ws_url_base = API_BASE_URL.replace("http", "ws")
|
125 |
ws_url = f"{ws_url_base}/ws/{token}"
|
@@ -128,76 +98,62 @@ async def listen_to_websockets(token: str, notification_state: list):
|
|
128 |
try:
|
129 |
async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
|
130 |
logger.info(f"[{ws_listener_id}] WebSocket connected successfully to {ws_url}")
|
131 |
-
|
132 |
-
# --- REMOVE or COMMENT OUT this debug line ---
|
133 |
-
# await asyncio.sleep(0.5)
|
134 |
-
# logger.info(f"[{ws_listener_id}] Connections according to manager after connect: {manager.active_connections}")
|
135 |
-
# --- End removal ---
|
136 |
-
|
137 |
while True:
|
138 |
try:
|
139 |
-
message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0)
|
140 |
logger.info(f"[{ws_listener_id}] Received raw message: {message_str}")
|
141 |
try:
|
142 |
message_data = json.loads(message_str)
|
143 |
logger.info(f"[{ws_listener_id}] Parsed message data: {message_data}")
|
144 |
|
145 |
if message_data.get("type") == "new_user":
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
else:
|
153 |
logger.warning(f"[{ws_listener_id}] Received message of unknown type: {message_data.get('type')}")
|
154 |
-
|
155 |
-
except json.JSONDecodeError:
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
except
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
logger.error(f"[{ws_listener_id}] WebSocket connection closed with error: {e}")
|
168 |
-
break
|
169 |
-
except Exception as e:
|
170 |
-
logger.error(f"[{ws_listener_id}] Error in WebSocket listener receive loop: {e}")
|
171 |
-
await asyncio.sleep(1)
|
172 |
-
|
173 |
-
except asyncio.TimeoutError: # Catch timeout during initial connect
|
174 |
-
logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out: {ws_url}")
|
175 |
-
except websockets.exceptions.InvalidURI:
|
176 |
-
logger.error(f"[{ws_listener_id}] Invalid WebSocket URI: {ws_url}")
|
177 |
-
except websockets.exceptions.WebSocketException as e: # Catch connection errors
|
178 |
-
logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
|
179 |
-
except Exception as e:
|
180 |
-
logger.error(f"[{ws_listener_id}] Unexpected error in WebSocket listener task: {e}")
|
181 |
|
182 |
logger.info(f"[{ws_listener_id}] Listener task finished.")
|
183 |
-
|
184 |
-
|
185 |
|
186 |
# --- Gradio Interface ---
|
187 |
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
188 |
# State variables
|
189 |
-
# Holds the session token after login
|
190 |
auth_token = gr.State(None)
|
191 |
-
# Holds user info {id, email} after login
|
192 |
user_info = gr.State(None)
|
193 |
-
# Holds the list of notification messages
|
194 |
notification_list = gr.State([])
|
195 |
-
# Holds the asyncio task for the WebSocket listener
|
196 |
websocket_task = gr.State(None)
|
|
|
|
|
197 |
|
198 |
# --- UI Components ---
|
199 |
with gr.Tabs() as tabs:
|
200 |
-
# --- Registration
|
201 |
with gr.TabItem("Register", id="register_tab"):
|
202 |
gr.Markdown("## Create a new account")
|
203 |
reg_email = gr.Textbox(label="Email", type="email")
|
@@ -205,8 +161,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
|
205 |
reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
|
206 |
reg_button = gr.Button("Register")
|
207 |
reg_status = gr.Textbox(label="Status", interactive=False)
|
208 |
-
|
209 |
-
# --- Login Tab ---
|
210 |
with gr.TabItem("Login", id="login_tab"):
|
211 |
gr.Markdown("## Login to your account")
|
212 |
login_email = gr.Textbox(label="Email", type="email")
|
@@ -214,172 +168,143 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
|
214 |
login_button = gr.Button("Login")
|
215 |
login_status = gr.Textbox(label="Status", interactive=False)
|
216 |
|
217 |
-
# --- Welcome Tab
|
218 |
with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
|
219 |
gr.Markdown("## Welcome!", elem_id="welcome_header")
|
220 |
welcome_message = gr.Markdown("", elem_id="welcome_message")
|
221 |
logout_button = gr.Button("Logout")
|
222 |
-
gr.Markdown("---")
|
223 |
gr.Markdown("## Real-time Notifications")
|
224 |
-
# Textbox to display notifications, updated periodically
|
225 |
notification_display = gr.Textbox(
|
226 |
label="New User Alerts",
|
227 |
lines=5,
|
228 |
max_lines=10,
|
229 |
interactive=False,
|
230 |
-
#
|
231 |
-
# This function will read the `notification_list` state
|
232 |
-
every=1
|
233 |
)
|
|
|
|
|
234 |
|
235 |
-
def update_notification_ui(notif_list_state):
|
236 |
-
# notif_list_state here *is* the Python list from the gr.State object
|
237 |
-
log_msg = f"UI Update Triggered. State List Length: {len(notif_list_state)}. Content: {notif_list_state[:5]}"
|
238 |
-
logger.debug(log_msg) # Keep logging
|
239 |
-
# --- Ensure we return an update object ---
|
240 |
-
# Join the list items into a string for display
|
241 |
-
new_value = "\n".join(notif_list_state)
|
242 |
-
return gr.update(value=new_value) # Explicitly return gr.update()
|
243 |
-
|
244 |
-
notification_display.change( # Use .change with every= setup on the component
|
245 |
-
fn=update_notification_ui,
|
246 |
-
inputs=[notification_list], # Read the state
|
247 |
-
outputs=[notification_display] # Update the component
|
248 |
-
)
|
249 |
|
250 |
# --- Event Handlers ---
|
251 |
|
252 |
-
# Registration Logic
|
253 |
async def handle_register(email, password, confirm_password):
|
254 |
-
if not email or not password or not confirm_password:
|
255 |
-
|
256 |
-
if password
|
257 |
-
return gr.update(value="Passwords do not match.")
|
258 |
-
if len(password) < 8:
|
259 |
-
return gr.update(value="Password must be at least 8 characters long.")
|
260 |
-
|
261 |
payload = {"email": email, "password": password}
|
262 |
result = await make_api_request("post", "/register", json=payload)
|
|
|
|
|
263 |
|
264 |
-
|
265 |
-
return gr.update(value=f"Registration failed: {result['error']}")
|
266 |
-
else:
|
267 |
-
# Optionally switch to login tab after successful registration
|
268 |
-
return gr.update(value=f"Registration successful for {result.get('email')}! Please log in.")
|
269 |
-
|
270 |
-
reg_button.click(
|
271 |
-
handle_register,
|
272 |
-
inputs=[reg_email, reg_password, reg_confirm_password],
|
273 |
-
outputs=[reg_status]
|
274 |
-
)
|
275 |
|
276 |
# Login Logic
|
277 |
-
|
278 |
-
|
279 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
280 |
|
281 |
payload = {"email": email, "password": password}
|
282 |
result = await make_api_request("post", "/login", json=payload)
|
283 |
|
284 |
if "error" in result:
|
285 |
-
|
|
|
286 |
else:
|
287 |
token = result.get("access_token")
|
288 |
-
|
289 |
-
user_data = await dependencies.get_optional_current_user(token) # Use dependency directly
|
290 |
-
|
291 |
if not user_data:
|
292 |
-
|
293 |
-
return
|
294 |
|
295 |
-
# Cancel any existing websocket listener task before starting a new one
|
296 |
if current_task and not current_task.done():
|
297 |
current_task.cancel()
|
298 |
-
try:
|
299 |
-
|
300 |
-
except asyncio.CancelledError:
|
301 |
-
logger.info("Previous WebSocket task cancelled.")
|
302 |
|
303 |
-
#
|
304 |
-
|
305 |
-
new_task = asyncio.create_task(listen_to_websockets(token, notification_list.value)) # Pass the list
|
306 |
|
307 |
-
# Update state and UI
|
308 |
welcome_msg = f"Welcome, {user_data.email}!"
|
309 |
-
#
|
310 |
return (
|
311 |
gr.update(value="Login successful!"), # login_status
|
312 |
token, # auth_token state
|
313 |
-
user_data.model_dump(), # user_info state
|
314 |
gr.update(selected="welcome_tab"), # Switch Tabs
|
315 |
gr.update(visible=True), # Make welcome tab visible
|
316 |
gr.update(value=welcome_msg), # Update welcome message markdown
|
317 |
-
new_task
|
|
|
318 |
)
|
319 |
|
|
|
320 |
login_button.click(
|
321 |
handle_login,
|
322 |
-
inputs=[login_email, login_password, websocket_task],
|
323 |
-
outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task]
|
324 |
)
|
325 |
|
326 |
-
|
327 |
-
#
|
328 |
-
|
329 |
-
|
330 |
-
#
|
331 |
-
|
332 |
-
|
333 |
-
|
|
|
|
|
|
|
|
|
|
|
334 |
fn=update_notification_ui,
|
335 |
-
inputs=[notification_list],
|
336 |
-
outputs=[notification_display]
|
337 |
)
|
338 |
|
|
|
|
|
|
|
|
|
339 |
|
340 |
# Logout Logic
|
|
|
341 |
async def handle_logout(current_task):
|
342 |
-
# Cancel the websocket listener task if it's running
|
343 |
if current_task and not current_task.done():
|
344 |
current_task.cancel()
|
345 |
-
try:
|
346 |
-
|
347 |
-
|
348 |
-
|
349 |
-
|
350 |
-
|
351 |
-
|
352 |
-
|
353 |
-
|
354 |
-
[], # Clear notifications
|
355 |
-
None, # Clear websocket_task
|
356 |
-
gr.update(selected="login_tab"),# Switch Tabs
|
357 |
-
gr.update(visible=False), # Hide welcome tab
|
358 |
-
gr.update(value=""), # Clear welcome message
|
359 |
-
gr.update(value="") # Clear login status
|
360 |
-
)
|
361 |
-
|
362 |
logout_button.click(
|
363 |
handle_logout,
|
364 |
inputs=[websocket_task],
|
365 |
outputs=[
|
366 |
-
auth_token,
|
367 |
-
|
368 |
-
|
369 |
-
websocket_task,
|
370 |
-
tabs,
|
371 |
-
welcome_tab,
|
372 |
-
welcome_message,
|
373 |
-
login_status
|
374 |
]
|
375 |
)
|
376 |
|
377 |
-
# Mount
|
378 |
app = gr.mount_gradio_app(app, demo, path="/")
|
379 |
|
380 |
-
#
|
381 |
-
# Use uvicorn to run the FastAPI app (which now includes Gradio)
|
382 |
if __name__ == "__main__":
|
383 |
import uvicorn
|
384 |
-
|
385 |
-
uvicorn.run(app, host="0.0.0.0", port=7860)
|
|
|
9 |
from contextlib import asynccontextmanager
|
10 |
|
11 |
from fastapi import FastAPI, Depends # Import FastAPI itself
|
|
|
12 |
from .database import connect_db, disconnect_db, database, metadata, users
|
13 |
from .api import router as api_router
|
14 |
from . import schemas, auth, dependencies
|
15 |
+
from .websocket import manager # Import the connection manager instance
|
16 |
+
|
17 |
from sqlalchemy.schema import CreateTable
|
18 |
+
from sqlalchemy.dialects import sqlite
|
19 |
|
|
|
20 |
logging.basicConfig(level=logging.INFO)
|
21 |
logger = logging.getLogger(__name__)
|
22 |
|
23 |
+
API_BASE_URL = "http://127.0.0.1:7860/api"
|
|
|
24 |
|
25 |
+
# --- Lifespan (remains the same) ---
|
26 |
@asynccontextmanager
|
27 |
async def lifespan(app: FastAPI):
|
28 |
+
# ... (same DB setup code) ...
|
29 |
logger.info("Application startup: Connecting DB...")
|
30 |
await connect_db()
|
31 |
logger.info("Application startup: DB Connected. Checking/Creating tables...")
|
32 |
+
if database.is_connected:
|
33 |
try:
|
|
|
|
|
34 |
check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;"
|
35 |
table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name})
|
|
|
36 |
if not table_exists:
|
37 |
logger.info(f"Table '{users.name}' not found, attempting creation using async connection...")
|
38 |
+
dialect = sqlite.dialect()
|
|
|
|
|
|
|
|
|
39 |
create_table_stmt = str(CreateTable(users).compile(dialect=dialect))
|
|
|
|
|
|
|
40 |
await database.execute(query=create_table_stmt)
|
41 |
logger.info(f"Table '{users.name}' created successfully via async connection.")
|
|
|
|
|
42 |
table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name})
|
43 |
+
if table_exists_after: logger.info(f"Table '{users.name}' verified after creation.")
|
44 |
+
else: logger.error(f"Table '{users.name}' verification FAILED after creation attempt!")
|
|
|
|
|
|
|
45 |
else:
|
46 |
logger.info(f"Table '{users.name}' already exists (checked via async connection).")
|
|
|
47 |
except Exception as db_setup_err:
|
48 |
logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}")
|
|
|
49 |
else:
|
50 |
logger.error("CRITICAL: Database connection failed, skipping table setup.")
|
|
|
51 |
logger.info("Application startup: DB setup phase complete.")
|
52 |
yield
|
|
|
53 |
logger.info("Application shutdown: Disconnecting DB...")
|
54 |
await disconnect_db()
|
55 |
logger.info("Application shutdown: DB Disconnected.")
|
56 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
57 |
|
58 |
+
# --- FastAPI App Setup (remains the same) ---
|
59 |
+
app = FastAPI(lifespan=lifespan)
|
60 |
+
app.include_router(api_router, prefix="/api")
|
61 |
|
62 |
+
# --- Helper functions (make_api_request remains the same) ---
|
63 |
async def make_api_request(method: str, endpoint: str, **kwargs):
|
64 |
async with httpx.AsyncClient() as client:
|
65 |
url = f"{API_BASE_URL}{endpoint}"
|
66 |
try:
|
67 |
response = await client.request(method, url, **kwargs)
|
68 |
+
response.raise_for_status()
|
69 |
return response.json()
|
70 |
except httpx.RequestError as e:
|
71 |
logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}")
|
72 |
return {"error": f"Network error contacting API: {e}"}
|
73 |
except httpx.HTTPStatusError as e:
|
74 |
logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
|
75 |
+
try: detail = e.response.json().get("detail", e.response.text)
|
76 |
+
except json.JSONDecodeError: detail = e.response.text
|
|
|
|
|
77 |
return {"error": f"API Error: {detail}"}
|
78 |
except Exception as e:
|
79 |
logger.error(f"Unexpected error during API call: {e}")
|
80 |
return {"error": f"An unexpected error occurred: {str(e)}"}
|
81 |
|
82 |
# --- WebSocket handling within Gradio ---
|
83 |
+
# <<< MODIFIED: Accept trigger state object >>>
|
84 |
+
async def listen_to_websockets(token: str, notification_list_state: gr.State, notification_trigger_state: gr.State):
|
85 |
+
"""Connects to WS and updates state list and trigger when a message arrives."""
|
86 |
ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
|
87 |
logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")
|
88 |
|
89 |
if not token:
|
90 |
logger.warning(f"[{ws_listener_id}] No token provided. Listener task exiting.")
|
91 |
+
# <<< Return original state values >>>
|
92 |
+
return notification_list_state.value, notification_trigger_state.value
|
93 |
|
94 |
ws_url_base = API_BASE_URL.replace("http", "ws")
|
95 |
ws_url = f"{ws_url_base}/ws/{token}"
|
|
|
98 |
try:
|
99 |
async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
|
100 |
logger.info(f"[{ws_listener_id}] WebSocket connected successfully to {ws_url}")
|
|
|
|
|
|
|
|
|
|
|
|
|
101 |
while True:
|
102 |
try:
|
103 |
+
message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0)
|
104 |
logger.info(f"[{ws_listener_id}] Received raw message: {message_str}")
|
105 |
try:
|
106 |
message_data = json.loads(message_str)
|
107 |
logger.info(f"[{ws_listener_id}] Parsed message data: {message_data}")
|
108 |
|
109 |
if message_data.get("type") == "new_user":
|
110 |
+
notification = schemas.Notification(**message_data)
|
111 |
+
logger.info(f"[{ws_listener_id}] Processing 'new_user' notification: {notification.message}")
|
112 |
+
|
113 |
+
# <<< Modify state objects' values >>>
|
114 |
+
current_list = notification_list_state.value.copy()
|
115 |
+
current_list.insert(0, notification.message)
|
116 |
+
if len(current_list) > 10: current_list.pop()
|
117 |
+
notification_list_state.value = current_list
|
118 |
+
logger.info(f"[{ws_listener_id}] State list updated via state object. New length: {len(notification_list_state.value)}. Content: {notification_list_state.value[:5]}")
|
119 |
+
|
120 |
+
# <<< Update trigger state object's value >>>
|
121 |
+
notification_trigger_state.value += 1
|
122 |
+
logger.info(f"[{ws_listener_id}] Incremented notification trigger to {notification_trigger_state.value}")
|
123 |
+
|
124 |
else:
|
125 |
logger.warning(f"[{ws_listener_id}] Received message of unknown type: {message_data.get('type')}")
|
126 |
+
# ... (error handling for parsing) ...
|
127 |
+
except json.JSONDecodeError: logger.error(f"[{ws_listener_id}] Failed to decode JSON: {message_str}")
|
128 |
+
except Exception as parse_err: logger.error(f"[{ws_listener_id}] Error processing message: {parse_err}")
|
129 |
+
# ... (error handling for websocket recv/connection) ...
|
130 |
+
except asyncio.TimeoutError: logger.debug(f"[{ws_listener_id}] WebSocket recv timed out."); continue
|
131 |
+
except websockets.ConnectionClosedOK: logger.info(f"[{ws_listener_id}] WebSocket connection closed normally."); break
|
132 |
+
except websockets.ConnectionClosedError as e: logger.error(f"[{ws_listener_id}] WebSocket connection closed with error: {e}"); break
|
133 |
+
except Exception as e: logger.error(f"[{ws_listener_id}] Error in listener receive loop: {e}"); await asyncio.sleep(1); break # Break on unknown errors too
|
134 |
+
# ... (error handling for websocket connect) ...
|
135 |
+
except asyncio.TimeoutError: logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out: {ws_url}")
|
136 |
+
except websockets.exceptions.InvalidURI: logger.error(f"[{ws_listener_id}] Invalid WebSocket URI: {ws_url}")
|
137 |
+
except websockets.exceptions.WebSocketException as e: logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
|
138 |
+
except Exception as e: logger.error(f"[{ws_listener_id}] Unexpected error in WebSocket listener task: {e}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
139 |
|
140 |
logger.info(f"[{ws_listener_id}] Listener task finished.")
|
141 |
+
# <<< Return original state values >>>
|
142 |
+
return notification_list_state.value, notification_trigger_state.value
|
143 |
|
144 |
# --- Gradio Interface ---
|
145 |
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
146 |
# State variables
|
|
|
147 |
auth_token = gr.State(None)
|
|
|
148 |
user_info = gr.State(None)
|
|
|
149 |
notification_list = gr.State([])
|
|
|
150 |
websocket_task = gr.State(None)
|
151 |
+
# <<< Add Dummy State >>>
|
152 |
+
notification_trigger = gr.State(0) # Simple counter
|
153 |
|
154 |
# --- UI Components ---
|
155 |
with gr.Tabs() as tabs:
|
156 |
+
# --- Registration/Login Tabs (remain the same) ---
|
157 |
with gr.TabItem("Register", id="register_tab"):
|
158 |
gr.Markdown("## Create a new account")
|
159 |
reg_email = gr.Textbox(label="Email", type="email")
|
|
|
161 |
reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
|
162 |
reg_button = gr.Button("Register")
|
163 |
reg_status = gr.Textbox(label="Status", interactive=False)
|
|
|
|
|
164 |
with gr.TabItem("Login", id="login_tab"):
|
165 |
gr.Markdown("## Login to your account")
|
166 |
login_email = gr.Textbox(label="Email", type="email")
|
|
|
168 |
login_button = gr.Button("Login")
|
169 |
login_status = gr.Textbox(label="Status", interactive=False)
|
170 |
|
171 |
+
# --- Welcome Tab ---
|
172 |
with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
|
173 |
gr.Markdown("## Welcome!", elem_id="welcome_header")
|
174 |
welcome_message = gr.Markdown("", elem_id="welcome_message")
|
175 |
logout_button = gr.Button("Logout")
|
176 |
+
gr.Markdown("---")
|
177 |
gr.Markdown("## Real-time Notifications")
|
|
|
178 |
notification_display = gr.Textbox(
|
179 |
label="New User Alerts",
|
180 |
lines=5,
|
181 |
max_lines=10,
|
182 |
interactive=False,
|
183 |
+
# <<< REMOVE every=1 >>>
|
|
|
|
|
184 |
)
|
185 |
+
# <<< Add Dummy Component (Hidden) >>>
|
186 |
+
dummy_trigger_output = gr.Textbox(label="trigger", visible=False)
|
187 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
188 |
|
189 |
# --- Event Handlers ---
|
190 |
|
191 |
+
# Registration Logic (remains the same)
|
192 |
async def handle_register(email, password, confirm_password):
|
193 |
+
if not email or not password or not confirm_password: return gr.update(value="Please fill in all fields.")
|
194 |
+
if password != confirm_password: return gr.update(value="Passwords do not match.")
|
195 |
+
if len(password) < 8: return gr.update(value="Password must be at least 8 characters long.")
|
|
|
|
|
|
|
|
|
196 |
payload = {"email": email, "password": password}
|
197 |
result = await make_api_request("post", "/register", json=payload)
|
198 |
+
if "error" in result: return gr.update(value=f"Registration failed: {result['error']}")
|
199 |
+
else: return gr.update(value=f"Registration successful for {result.get('email')}! Please log in.")
|
200 |
|
201 |
+
reg_button.click(handle_register, inputs=[reg_email, reg_password, reg_confirm_password], outputs=[reg_status])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
202 |
|
203 |
# Login Logic
|
204 |
+
# <<< MODIFIED: Pass trigger state, update outputs >>>
|
205 |
+
async def handle_login(email, password, current_task, current_trigger_val):
|
206 |
+
# --- Output state needs to align with return values ---
|
207 |
+
# Returns: login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger
|
208 |
+
outputs_tuple = (
|
209 |
+
gr.update(value="Please enter email and password."), # login_status
|
210 |
+
None, None, None, gr.update(visible=False), None, # auth_token, user_info, tabs, welcome_tab, welcome_message
|
211 |
+
current_task, # websocket_task (no change)
|
212 |
+
current_trigger_val # notification_trigger (no change)
|
213 |
+
)
|
214 |
+
if not email or not password: return outputs_tuple
|
215 |
|
216 |
payload = {"email": email, "password": password}
|
217 |
result = await make_api_request("post", "/login", json=payload)
|
218 |
|
219 |
if "error" in result:
|
220 |
+
outputs_tuple = (gr.update(value=f"Login failed: {result['error']}"), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val)
|
221 |
+
return outputs_tuple
|
222 |
else:
|
223 |
token = result.get("access_token")
|
224 |
+
user_data = await dependencies.get_optional_current_user(token)
|
|
|
|
|
225 |
if not user_data:
|
226 |
+
outputs_tuple = (gr.update(value="Login succeeded but failed to fetch user data."), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val)
|
227 |
+
return outputs_tuple
|
228 |
|
|
|
229 |
if current_task and not current_task.done():
|
230 |
current_task.cancel()
|
231 |
+
try: await current_task
|
232 |
+
except asyncio.CancelledError: logger.info("Previous WebSocket task cancelled.")
|
|
|
|
|
233 |
|
234 |
+
# <<< Pass both state objects to listener >>>
|
235 |
+
new_task = asyncio.create_task(listen_to_websockets(token, notification_list, notification_trigger))
|
|
|
236 |
|
|
|
237 |
welcome_msg = f"Welcome, {user_data.email}!"
|
238 |
+
# --- Ensure number of return values matches outputs ---
|
239 |
return (
|
240 |
gr.update(value="Login successful!"), # login_status
|
241 |
token, # auth_token state
|
242 |
+
user_data.model_dump(), # user_info state
|
243 |
gr.update(selected="welcome_tab"), # Switch Tabs
|
244 |
gr.update(visible=True), # Make welcome tab visible
|
245 |
gr.update(value=welcome_msg), # Update welcome message markdown
|
246 |
+
new_task, # websocket_task state
|
247 |
+
0 # Reset notification_trigger state to 0 on login
|
248 |
)
|
249 |
|
250 |
+
# <<< MODIFIED: Add notification_trigger to inputs/outputs >>>
|
251 |
login_button.click(
|
252 |
handle_login,
|
253 |
+
inputs=[login_email, login_password, websocket_task, notification_trigger],
|
254 |
+
outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger]
|
255 |
)
|
256 |
|
257 |
+
# Function to update the notification display based on the list state
|
258 |
+
# <<< Triggered by dummy component now >>>
|
259 |
+
def update_notification_ui(notif_list): # Takes the list value directly now
|
260 |
+
log_msg = f"UI Update Triggered via Dummy. State List Length: {len(notif_list)}. Content: {notif_list[:5]}"
|
261 |
+
logger.info(log_msg) # Use info level to ensure visibility
|
262 |
+
new_value = "\n".join(notif_list)
|
263 |
+
return gr.update(value=new_value)
|
264 |
+
|
265 |
+
# <<< Add Event handler for the dummy trigger >>>
|
266 |
+
# When the dummy_trigger_output *would* change (because notification_trigger state changed)...
|
267 |
+
# ...call the update_notification_ui function.
|
268 |
+
# Pass the *current* value of notification_list state as input to the function.
|
269 |
+
dummy_trigger_output.change(
|
270 |
fn=update_notification_ui,
|
271 |
+
inputs=[notification_list], # Input is the list state we want to display
|
272 |
+
outputs=[notification_display] # Output updates the real display textbox
|
273 |
)
|
274 |
|
275 |
+
# <<< Link the dummy trigger state to the dummy output component >>>
|
276 |
+
# This makes the dummy_trigger_output.change event fire when notification_trigger changes
|
277 |
+
notification_trigger.change(lambda x: x, inputs=notification_trigger, outputs=dummy_trigger_output, queue=False) # Use queue=False for immediate trigger if possible
|
278 |
+
|
279 |
|
280 |
# Logout Logic
|
281 |
+
# <<< MODIFIED: Add notification_trigger to outputs >>>
|
282 |
async def handle_logout(current_task):
|
|
|
283 |
if current_task and not current_task.done():
|
284 |
current_task.cancel()
|
285 |
+
try: await current_task
|
286 |
+
except asyncio.CancelledError: logger.info("WebSocket task cancelled on logout.")
|
287 |
+
# --- Ensure number of return values matches outputs ---
|
288 |
+
return ( None, None, [], None, # auth_token, user_info, notification_list, websocket_task
|
289 |
+
gr.update(selected="login_tab"), gr.update(visible=False), # tabs, welcome_tab
|
290 |
+
gr.update(value=""), gr.update(value=""), # welcome_message, login_status
|
291 |
+
0 ) # notification_trigger (reset to 0)
|
292 |
+
|
293 |
+
# <<< MODIFIED: Add notification_trigger to outputs >>>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
294 |
logout_button.click(
|
295 |
handle_logout,
|
296 |
inputs=[websocket_task],
|
297 |
outputs=[
|
298 |
+
auth_token, user_info, notification_list, websocket_task,
|
299 |
+
tabs, welcome_tab, welcome_message, login_status,
|
300 |
+
notification_trigger # Add trigger to outputs
|
|
|
|
|
|
|
|
|
|
|
301 |
]
|
302 |
)
|
303 |
|
304 |
+
# Mount Gradio App (remains the same)
|
305 |
app = gr.mount_gradio_app(app, demo, path="/")
|
306 |
|
307 |
+
# Run Uvicorn (remains the same)
|
|
|
308 |
if __name__ == "__main__":
|
309 |
import uvicorn
|
310 |
+
uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=True)
|
|