amaye15 commited on
Commit
3222a21
·
1 Parent(s): fbb85b5
Files changed (4) hide show
  1. Dockerfile +7 -11
  2. app/database.py +138 -64
  3. app/main.py +256 -263
  4. requirements.txt +10 -7
Dockerfile CHANGED
@@ -11,18 +11,14 @@ COPY ./requirements.txt /code/requirements.txt
11
  RUN pip install --no-cache-dir --upgrade pip && \
12
  pip install --no-cache-dir -r requirements.txt
13
 
14
- # Copy the rest of the application code into the container
15
  COPY ./app /code/app
16
- # COPY ./.env /code/.env # Copy .env file - for Hugging Face, use Secrets instead
17
 
18
- # Make port 7860 available to the world outside this container (Gradio default)
19
  EXPOSE 7860
20
 
21
- # Explicitly create the /data directory where the SQLite DB will live
22
- # Running as root by default, so permissions should be okay initially
23
- # RUN mkdir -p /data
24
-
25
- # Command to run the application using uvicorn
26
- # It will run the FastAPI app instance created in app/main.py
27
- # Host 0.0.0.0 is important to accept connections from outside the container
28
- CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
 
11
  RUN pip install --no-cache-dir --upgrade pip && \
12
  pip install --no-cache-dir -r requirements.txt
13
 
14
+ # Copy the application code (Backend API and Streamlit app)
15
  COPY ./app /code/app
16
+ COPY ./streamlit_app.py /code/streamlit_app.py # Add streamlit app file
17
 
18
+ # Make port 7860 available (Streamlit default is 8501, but HF uses specified port)
19
  EXPOSE 7860
20
 
21
+ # Command to run the Streamlit application
22
+ # Use --server.port to match EXPOSE and HF config
23
+ # Use --server.address to bind correctly inside container
24
+ CMD ["streamlit", "run", "streamlit_app.py", "--server.port=7860", "--server.address=0.0.0.0"]
 
 
 
 
app/database.py CHANGED
@@ -1,37 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # app/database.py
2
  import os
3
- from databases import Database
 
4
  from dotenv import load_dotenv
5
- # --- Keep only these SQLAlchemy imports ---
6
- from sqlalchemy import MetaData, Table, Column, Integer, String
7
  import logging
8
- from urllib.parse import urlparse, urlunparse, parse_qs, urlencode
9
 
10
  load_dotenv()
11
  logger = logging.getLogger(__name__)
12
 
13
- # --- Database URL Configuration ---
14
- DEFAULT_DB_PATH = "/tmp/app.db" # Store DB in the temporary directory
15
- raw_db_url = os.getenv("DATABASE_URL", f"sqlite+aiosqlite:///{DEFAULT_DB_PATH}")
16
-
17
- final_database_url = raw_db_url
18
- if raw_db_url.startswith("sqlite+aiosqlite"):
19
- parsed_url = urlparse(raw_db_url)
20
- query_params = parse_qs(parsed_url.query)
21
- if 'check_same_thread' not in query_params:
22
- query_params['check_same_thread'] = ['False']
23
- new_query = urlencode(query_params, doseq=True)
24
- final_database_url = urlunparse(parsed_url._replace(query=new_query))
25
- logger.info(f"Using final async DB URL: {final_database_url}")
26
- else:
27
- logger.info(f"Using non-SQLite async DB URL: {final_database_url}")
28
-
29
- # --- Async Database Instance ---
30
- database = Database(final_database_url)
31
-
32
- # --- Metadata and Table Definition (Still needed for DDL generation) ---
33
  metadata = MetaData()
34
- users = Table(
35
  "users",
36
  metadata,
37
  Column("id", Integer, primary_key=True),
@@ -39,46 +114,45 @@ users = Table(
39
  Column("hashed_password", String, nullable=False),
40
  )
41
 
42
- # --- REMOVE ALL SYNCHRONOUS ENGINE AND TABLE CREATION LOGIC ---
43
-
44
- # --- Keep and refine Async connect/disconnect functions ---
45
- async def connect_db():
46
- """Connects to the database, ensuring the parent directory exists."""
47
  try:
48
- # Ensure the directory exists just before connecting
49
- db_file_path = final_database_url.split("sqlite:///")[-1].split("?")[0]
50
  db_dir = os.path.dirname(db_file_path)
51
- if db_dir: # Only proceed if a directory path was found
52
- if not os.path.exists(db_dir):
53
- logger.info(f"Database directory {db_dir} does not exist. Attempting creation...")
54
- try:
55
- os.makedirs(db_dir, exist_ok=True)
56
- logger.info(f"Created database directory {db_dir}.")
57
- except Exception as mkdir_err:
58
- # Log error but proceed, connection might still work if path is valid but dir creation failed weirdly
59
- logger.error(f"Failed to create directory {db_dir}: {mkdir_err}")
60
- # Check writability after ensuring existence attempt
61
- if os.path.exists(db_dir) and not os.access(db_dir, os.W_OK):
62
- logger.error(f"CRITICAL: Directory {db_dir} exists but is not writable!")
63
- elif not os.path.exists(db_dir):
64
- logger.error(f"CRITICAL: Directory {db_dir} does not exist and could not be created!")
65
-
66
-
67
- # Now attempt connection
68
- await database.connect()
69
- logger.info(f"Database connection established (async): {final_database_url}")
70
- # Table creation will happen in main.py lifespan event using this connection
71
- except Exception as e:
72
- logger.exception(f"Failed to establish async database connection: {e}")
73
- raise # Reraise critical error during startup
 
 
 
 
 
 
 
 
74
 
75
- async def disconnect_db():
76
- """Disconnects from the database if connected."""
77
- try:
78
- if database.is_connected:
79
- await database.disconnect()
80
- logger.info("Database connection closed (async).")
81
- else:
82
- logger.info("Database already disconnected (async).")
83
  except Exception as e:
84
- logger.exception(f"Error closing async database connection: {e}")
 
 
1
+ # # app/database.py
2
+ # import os
3
+ # from databases import Database
4
+ # from dotenv import load_dotenv
5
+ # # --- Keep only these SQLAlchemy imports ---
6
+ # from sqlalchemy import MetaData, Table, Column, Integer, String
7
+ # import logging
8
+ # from urllib.parse import urlparse, urlunparse, parse_qs, urlencode
9
+
10
+ # load_dotenv()
11
+ # logger = logging.getLogger(__name__)
12
+
13
+ # # --- Database URL Configuration ---
14
+ # DEFAULT_DB_PATH = "/tmp/app.db" # Store DB in the temporary directory
15
+ # raw_db_url = os.getenv("DATABASE_URL", f"sqlite+aiosqlite:///{DEFAULT_DB_PATH}")
16
+
17
+ # final_database_url = raw_db_url
18
+ # if raw_db_url.startswith("sqlite+aiosqlite"):
19
+ # parsed_url = urlparse(raw_db_url)
20
+ # query_params = parse_qs(parsed_url.query)
21
+ # if 'check_same_thread' not in query_params:
22
+ # query_params['check_same_thread'] = ['False']
23
+ # new_query = urlencode(query_params, doseq=True)
24
+ # final_database_url = urlunparse(parsed_url._replace(query=new_query))
25
+ # logger.info(f"Using final async DB URL: {final_database_url}")
26
+ # else:
27
+ # logger.info(f"Using non-SQLite async DB URL: {final_database_url}")
28
+
29
+ # # --- Async Database Instance ---
30
+ # database = Database(final_database_url)
31
+
32
+ # # --- Metadata and Table Definition (Still needed for DDL generation) ---
33
+ # metadata = MetaData()
34
+ # users = Table(
35
+ # "users",
36
+ # metadata,
37
+ # Column("id", Integer, primary_key=True),
38
+ # Column("email", String, unique=True, index=True, nullable=False),
39
+ # Column("hashed_password", String, nullable=False),
40
+ # )
41
+
42
+ # # --- REMOVE ALL SYNCHRONOUS ENGINE AND TABLE CREATION LOGIC ---
43
+
44
+ # # --- Keep and refine Async connect/disconnect functions ---
45
+ # async def connect_db():
46
+ # """Connects to the database, ensuring the parent directory exists."""
47
+ # try:
48
+ # # Ensure the directory exists just before connecting
49
+ # db_file_path = final_database_url.split("sqlite:///")[-1].split("?")[0]
50
+ # db_dir = os.path.dirname(db_file_path)
51
+ # if db_dir: # Only proceed if a directory path was found
52
+ # if not os.path.exists(db_dir):
53
+ # logger.info(f"Database directory {db_dir} does not exist. Attempting creation...")
54
+ # try:
55
+ # os.makedirs(db_dir, exist_ok=True)
56
+ # logger.info(f"Created database directory {db_dir}.")
57
+ # except Exception as mkdir_err:
58
+ # # Log error but proceed, connection might still work if path is valid but dir creation failed weirdly
59
+ # logger.error(f"Failed to create directory {db_dir}: {mkdir_err}")
60
+ # # Check writability after ensuring existence attempt
61
+ # if os.path.exists(db_dir) and not os.access(db_dir, os.W_OK):
62
+ # logger.error(f"CRITICAL: Directory {db_dir} exists but is not writable!")
63
+ # elif not os.path.exists(db_dir):
64
+ # logger.error(f"CRITICAL: Directory {db_dir} does not exist and could not be created!")
65
+
66
+
67
+ # # Now attempt connection
68
+ # await database.connect()
69
+ # logger.info(f"Database connection established (async): {final_database_url}")
70
+ # # Table creation will happen in main.py lifespan event using this connection
71
+ # except Exception as e:
72
+ # logger.exception(f"Failed to establish async database connection: {e}")
73
+ # raise # Reraise critical error during startup
74
+
75
+ # async def disconnect_db():
76
+ # """Disconnects from the database if connected."""
77
+ # try:
78
+ # if database.is_connected:
79
+ # await database.disconnect()
80
+ # logger.info("Database connection closed (async).")
81
+ # else:
82
+ # logger.info("Database already disconnected (async).")
83
+ # except Exception as e:
84
+ # logger.exception(f"Error closing async database connection: {e}")
85
+
86
+
87
  # app/database.py
88
  import os
89
+ # --- Keep only Sync SQLAlchemy ---
90
+ from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String, text, exc as sqlalchemy_exc
91
  from dotenv import load_dotenv
 
 
92
  import logging
 
93
 
94
  load_dotenv()
95
  logger = logging.getLogger(__name__)
96
 
97
+ # Use /tmp for ephemeral storage in HF Space
98
+ DEFAULT_DB_PATH = "/tmp/app.db"
99
+ # Construct sync URL directly
100
+ DATABASE_URL = os.getenv("DATABASE_URL_SYNC", f"sqlite:///{DEFAULT_DB_PATH}")
101
+
102
+ logger.info(f"Using DB URL for sync operations: {DATABASE_URL}")
103
+
104
+ # SQLite specific args for sync engine
105
+ connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}
106
+ engine = create_engine(DATABASE_URL, connect_args=connect_args, echo=False) # echo=True for debugging SQL
107
+
 
 
 
 
 
 
 
 
 
108
  metadata = MetaData()
109
+ users_table = Table( # Renamed slightly to avoid confusion with Pydantic model name
110
  "users",
111
  metadata,
112
  Column("id", Integer, primary_key=True),
 
114
  Column("hashed_password", String, nullable=False),
115
  )
116
 
117
+ def ensure_db_and_table_exist():
118
+ """Synchronously ensures DB file directory and users table exist."""
119
+ logger.info("Ensuring DB and table exist...")
 
 
120
  try:
121
+ # Ensure directory exists
122
+ db_file_path = DATABASE_URL.split("sqlite:///")[-1]
123
  db_dir = os.path.dirname(db_file_path)
124
+ if db_dir:
125
+ if not os.path.exists(db_dir):
126
+ logger.info(f"Creating DB directory: {db_dir}")
127
+ os.makedirs(db_dir, exist_ok=True)
128
+ if not os.access(db_dir, os.W_OK):
129
+ logger.error(f"CRITICAL: Directory {db_dir} not writable!")
130
+ return # Cannot proceed
131
+
132
+ # Check/Create table using the engine
133
+ with engine.connect() as connection:
134
+ try:
135
+ connection.execute(text("SELECT 1 FROM users LIMIT 1"))
136
+ logger.info("Users table already exists.")
137
+ except sqlalchemy_exc.OperationalError as e:
138
+ if "no such table" in str(e).lower():
139
+ logger.warning("Users table not found, creating...")
140
+ metadata.create_all(bind=connection) # Use connection
141
+ connection.commit()
142
+ logger.info("Users table created and committed.")
143
+ # Verify
144
+ try:
145
+ connection.execute(text("SELECT 1 FROM users LIMIT 1"))
146
+ logger.info("Users table verified post-creation.")
147
+ except Exception as verify_err:
148
+ logger.error(f"Verification failed after creating table: {verify_err}")
149
+ else:
150
+ logger.error(f"DB OperationalError checking table (not 'no such table'): {e}")
151
+ raise # Re-raise unexpected errors
152
+ except Exception as check_err:
153
+ logger.error(f"Unexpected error checking table: {check_err}")
154
+ raise # Re-raise unexpected errors
155
 
 
 
 
 
 
 
 
 
156
  except Exception as e:
157
+ logger.exception(f"CRITICAL error during DB setup: {e}")
158
+ # Potentially raise to halt app start?
app/main.py CHANGED
@@ -1,276 +1,269 @@
1
- # app/main.py
2
- import gradio as gr
 
3
  import httpx
4
- import websockets
5
  import asyncio
 
6
  import json
7
- import os
 
8
  import logging
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: response = await client.request(method, url, **kwargs); response.raise_for_status(); return response.json()
67
- except httpx.RequestError as e: logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}"); return {"error": f"Network error: {e}"}
68
- except httpx.HTTPStatusError as e:
69
- logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
70
- try: detail = e.response.json().get("detail", e.response.text)
71
- except json.JSONDecodeError: detail = e.response.text
72
- return {"error": f"API Error: {detail}"}
73
- except Exception as e: logger.error(f"Unexpected error during API call: {e}"); return {"error": f"Unexpected error: {str(e)}"}
74
-
75
- # --- WebSocket handling ---
76
- # <<< Pass state objects by reference >>>
77
- async def listen_to_websockets(token: str, notification_list_state: gr.State, notification_trigger_state: gr.State):
78
- """Connects to WS and updates state list and trigger when a message arrives."""
79
- ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
80
- logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")
81
-
82
- if not token:
83
- logger.warning(f"[{ws_listener_id}] No token provided. Task exiting.")
84
- return # Just exit, don't need to return state values
85
-
86
- ws_url_base = API_BASE_URL.replace("http", "ws")
87
- ws_url = f"{ws_url_base}/ws/{token}"
88
- logger.info(f"[{ws_listener_id}] Attempting to connect: {ws_url}")
89
 
90
  try:
91
- async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
92
- logger.info(f"[{ws_listener_id}] WebSocket connected successfully.")
93
- while True:
94
- try:
95
- message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0)
96
- logger.info(f"[{ws_listener_id}] Received: {message_str[:100]}...")
97
- try:
98
- message_data = json.loads(message_str)
99
- if message_data.get("type") == "new_user":
100
- notification = schemas.Notification(**message_data)
101
- logger.info(f"[{ws_listener_id}] Processing 'new_user': {notification.message}")
102
- # --- Modify state values directly ---
103
- current_list = notification_list_state.value.copy() # Operate on copy
104
- current_list.insert(0, notification.message)
105
- if len(current_list) > 10: current_list.pop()
106
- notification_list_state.value = current_list # Assign back modified copy
107
- notification_trigger_state.value += 1 # Increment trigger
108
- # --- Log the update ---
109
- logger.info(f"[{ws_listener_id}] States updated: list len={len(notification_list_state.value)}, trigger={notification_trigger_state.value}")
110
- else:
111
- logger.warning(f"[{ws_listener_id}] Unknown message type: {message_data.get('type')}")
112
- # ... (error handling for parsing) ...
113
- except json.JSONDecodeError: logger.error(f"[{ws_listener_id}] JSON Decode Error: {message_str}")
114
- except Exception as parse_err: logger.error(f"[{ws_listener_id}] Message Processing Error: {parse_err}")
115
- # ... (error handling for websocket recv/connection) ...
116
- except asyncio.TimeoutError: logger.debug(f"[{ws_listener_id}] WebSocket recv timed out."); continue
117
- except websockets.ConnectionClosedOK: logger.info(f"[{ws_listener_id}] WebSocket closed normally."); break
118
- except websockets.ConnectionClosedError as e: logger.error(f"[{ws_listener_id}] WebSocket closed with error: {e}"); break
119
- except Exception as e: logger.error(f"[{ws_listener_id}] Listener loop error: {e}"); await asyncio.sleep(1); break
120
- # ... (error handling for websocket connect) ...
121
- except asyncio.TimeoutError: logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out.")
122
- except websockets.exceptions.InvalidURI: logger.error(f"[{ws_listener_id}] Invalid WebSocket URI.")
123
- except websockets.exceptions.WebSocketException as e: logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
124
- except Exception as e: logger.error(f"[{ws_listener_id}] Unexpected error in listener task: {e}")
125
-
126
- logger.info(f"[{ws_listener_id}] Listener task finished.")
127
- # No need to return state values when task ends
128
-
129
- # --- Gradio Interface ---
130
- with gr.Blocks(theme=gr.themes.Soft()) as demo:
131
- # State variables
132
- auth_token = gr.State(None)
133
- user_info = gr.State(None)
134
- notification_list = gr.State([])
135
- websocket_task = gr.State(None)
136
- # Add trigger states
137
- notification_trigger = gr.State(0)
138
- last_polled_trigger = gr.State(0) # State to track last seen trigger
139
-
140
- # --- UI Components ---
141
- with gr.Tabs() as tabs:
142
- # --- Registration/Login Tabs (remain the same) ---
143
- with gr.TabItem("Register", id="register_tab"):
144
- gr.Markdown("## Create a new account")
145
- reg_email = gr.Textbox(label="Email", type="email")
146
- reg_password = gr.Textbox(label="Password (min 8 chars)", type="password")
147
- reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
148
- reg_button = gr.Button("Register")
149
- reg_status = gr.Textbox(label="Status", interactive=False)
150
- with gr.TabItem("Login", id="login_tab"):
151
- gr.Markdown("## Login to your account")
152
- login_email = gr.Textbox(label="Email", type="email")
153
- login_password = gr.Textbox(label="Password", type="password")
154
- login_button = gr.Button("Login")
155
- login_status = gr.Textbox(label="Status", interactive=False)
156
-
157
- # --- Welcome Tab ---
158
- with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
159
- gr.Markdown("## Welcome!", elem_id="welcome_header")
160
- welcome_message = gr.Markdown("", elem_id="welcome_message")
161
- logout_button = gr.Button("Logout")
162
- gr.Markdown("---")
163
- gr.Markdown("## Real-time Notifications")
164
- notification_display = gr.Textbox( # Visible display
165
- label="New User Alerts", lines=5, max_lines=10, interactive=False
166
- )
167
- # <<< Hidden component for polling >>>
168
- dummy_poller = gr.Number(label="poller", value=0, visible=False, every=1)
169
-
170
-
171
- # --- Event Handlers ---
172
-
173
- # Registration Logic (remains the same)
174
- async def handle_register(email, password, confirm_password):
175
- if not email or not password or not confirm_password: return gr.update(value="Please fill fields.")
176
- if password != confirm_password: return gr.update(value="Passwords mismatch.")
177
- if len(password) < 8: return gr.update(value="Password >= 8 chars.")
178
- payload = {"email": email, "password": password}
179
- result = await make_api_request("post", "/register", json=payload)
180
- if "error" in result: return gr.update(value=f"Register failed: {result['error']}")
181
- else: return gr.update(value=f"Register success: {result.get('email')}! Log in.")
182
- reg_button.click(handle_register, inputs=[reg_email, reg_password, reg_confirm_password], outputs=[reg_status])
183
-
184
- # Login Logic
185
- # <<< MODIFIED: Pass/Reset both trigger states >>>
186
- async def handle_login(email, password, current_task, current_trigger_val, current_last_poll_val):
187
- # Define failure outputs matching the outputs list
188
- fail_outputs = (gr.update(value="..."), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val, current_last_poll_val)
189
-
190
- if not email or not password: return fail_outputs[:1] + (gr.update(value="Enter email/password"),) + fail_outputs[2:]
191
-
192
- payload = {"email": email, "password": password}
193
- result = await make_api_request("post", "/login", json=payload)
194
-
195
- if "error" in result: return (gr.update(value=f"Login failed: {result['error']}"),) + fail_outputs[1:]
196
- else:
197
- token = result.get("access_token")
198
- user_data = await dependencies.get_optional_current_user(token)
199
- if not user_data: return (gr.update(value="Login ok but user fetch failed."),) + fail_outputs[1:]
200
-
201
- if current_task and not current_task.done():
202
- current_task.cancel()
203
- try: await current_task
204
- except asyncio.CancelledError: logger.info("Previous WebSocket task cancelled.")
205
-
206
- # <<< Pass state objects to listener >>>
207
- new_task = asyncio.create_task(listen_to_websockets(token, notification_list, notification_trigger))
208
-
209
- welcome_msg = f"Welcome, {user_data.email}!"
210
- # Reset triggers on successful login
211
- return (
212
- gr.update(value="Login successful!"), token, user_data.model_dump(), # status, token, user_info
213
- gr.update(selected="welcome_tab"), gr.update(visible=True), gr.update(value=welcome_msg), # UI changes
214
- new_task, 0, 0 # websocket_task, notification_trigger, last_polled_trigger
215
- )
216
-
217
- # <<< MODIFIED: Add last_polled_trigger to inputs/outputs >>>
218
- login_button.click(
219
- handle_login,
220
- inputs=[login_email, login_password, websocket_task, notification_trigger, last_polled_trigger],
221
- outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger, last_polled_trigger]
222
- )
223
-
224
-
225
- # <<< Polling function >>>
226
- def poll_and_update(current_trigger_value, last_known_trigger, current_notif_list):
227
- """ Checks trigger state change and updates UI if needed. """
228
- if current_trigger_value != last_known_trigger:
229
- logger.info(f"Polling detected trigger change ({last_known_trigger} -> {current_trigger_value}). Updating UI.")
230
- new_value = "\n".join(current_notif_list)
231
- # Return new display value AND update last_known_trigger state value
232
- return gr.update(value=new_value), current_trigger_value
233
- else:
234
- # No change, return NoUpdate for display AND existing last_known_trigger value
235
- return gr.NoUpdate(), last_known_trigger
236
-
237
- # <<< Attach polling function to the dummy component's change event >>>
238
- dummy_poller.change(
239
- fn=poll_and_update,
240
- inputs=[notification_trigger, last_polled_trigger, notification_list],
241
- outputs=[notification_display, last_polled_trigger], # Update display & last polled state
242
- queue=False # Try immediate update
243
- )
244
-
245
-
246
- # Logout Logic
247
- # <<< MODIFIED: Reset both trigger states >>>
248
- async def handle_logout(current_task):
249
- if current_task and not current_task.done():
250
- current_task.cancel()
251
- try: await current_task
252
- except asyncio.CancelledError: logger.info("WebSocket task cancelled on logout.")
253
- # Reset all relevant states
254
- return ( None, None, [], None, # auth_token, user_info, notification_list, websocket_task
255
- gr.update(selected="login_tab"), gr.update(visible=False), # tabs, welcome_tab
256
- gr.update(value=""), gr.update(value=""), # welcome_message, login_status
257
- 0, 0 ) # notification_trigger, last_polled_trigger (reset)
258
-
259
- # <<< MODIFIED: Add last_polled_trigger to outputs >>>
260
- logout_button.click(
261
- handle_logout,
262
- inputs=[websocket_task],
263
- outputs=[
264
- auth_token, user_info, notification_list, websocket_task,
265
- tabs, welcome_tab, welcome_message, login_status,
266
- notification_trigger, last_polled_trigger # Add trigger states here
267
- ]
268
- )
269
-
270
- # Mount Gradio App (remains the same)
271
- app = gr.mount_gradio_app(app, demo, path="/")
272
-
273
- # Run Uvicorn (remains the same)
274
- if __name__ == "__main__":
275
- import uvicorn
276
- uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=True)
 
1
+ # streamlit_app.py
2
+ import streamlit as st
3
+ from streamlit_autorefresh import st_autorefresh # For periodic refresh
4
  import httpx
 
5
  import asyncio
6
+ import websockets
7
  import json
8
+ import threading
9
+ import queue
10
  import logging
11
+ import time
12
+ import os
13
 
14
+ # Import backend components
15
+ from app import crud, models, schemas, auth, dependencies
16
+ from app.database import ensure_db_and_table_exist # Sync function
17
+ from app.websocket import manager # Import the manager instance
 
18
 
19
+ # FastAPI imports for mounting
20
+ from fastapi import FastAPI, Depends, HTTPException, status
21
+ from fastapi.routing import Mount
22
+ from fastapi.staticfiles import StaticFiles
23
+ from app.api import router as api_router # Import the specific API router
24
 
25
+ # --- Logging ---
26
  logging.basicConfig(level=logging.INFO)
27
  logger = logging.getLogger(__name__)
28
 
29
+ # --- Configuration ---
30
+ # Use environment variable or default for local vs deployed API endpoint
31
+ # Since we are mounting FastAPI within Streamlit for the HF Space deployment:
32
+ API_BASE_URL = "http://127.0.0.1:7860/api" # Calls within the same process
33
+ WS_BASE_URL = API_BASE_URL.replace("http", "ws")
34
+
35
+ # --- Ensure DB exists on first run ---
36
+ # This runs once per session/process start
37
+ ensure_db_and_table_exist()
38
+
39
+ # --- FastAPI Mounting Setup ---
40
+ # Create a FastAPI instance (separate from the Streamlit one)
41
+ # We won't run this directly with uvicorn, but Streamlit uses it internally
42
+ api_app = FastAPI(title="Backend API") # Can add lifespan if needed for API-specific setup later
43
+ api_app.include_router(api_router, prefix="/api")
44
+
45
+ # Mount the FastAPI app within Streamlit's internal Tornado server
46
+ # This requires monkey-patching or using available hooks if Streamlit allows.
47
+ # Simpler approach for HF Space: Run FastAPI separately is cleaner if possible.
48
+ # Reverting to the idea that for this HF Space demo, API calls will be internal HTTP requests.
49
+
50
+ # --- WebSocket Listener Thread ---
51
+ stop_event = threading.Event()
52
+ notification_queue = queue.Queue()
53
+
54
+ def websocket_listener(token: str):
55
+ """Runs in a background thread to listen for WebSocket messages."""
56
+ logger.info(f"[WS Thread] Listener started for token: {token[:10]}...")
57
+ ws_url = f"{WS_BASE_URL}/ws/{token}"
58
+
59
+ async def listen():
60
  try:
61
+ async with websockets.connect(ws_url, open_timeout=10.0) as ws:
62
+ logger.info(f"[WS Thread] Connected to {ws_url}")
63
+ st.session_state['ws_connected'] = True
64
+ while not stop_event.is_set():
65
+ try:
66
+ message = await asyncio.wait_for(ws.recv(), timeout=1.0) # Check stop_event frequently
67
+ logger.info(f"[WS Thread] Received message: {message[:100]}...")
68
+ try:
69
+ data = json.loads(message)
70
+ if data.get("type") == "new_user":
71
+ notification = schemas.Notification(**data)
72
+ notification_queue.put(notification.message) # Put message in queue
73
+ logger.info("[WS Thread] Put notification in queue.")
74
+ except json.JSONDecodeError:
75
+ logger.error("[WS Thread] Failed to decode JSON.")
76
+ except Exception as e:
77
+ logger.error(f"[WS Thread] Error processing message: {e}")
78
+
79
+ except asyncio.TimeoutError:
80
+ continue # No message, check stop_event again
81
+ except websockets.ConnectionClosed:
82
+ logger.warning("[WS Thread] Connection closed.")
83
+ break # Exit loop if closed
84
+ except Exception as e:
85
+ logger.error(f"[WS Thread] Connection failed or error: {e}")
86
+ finally:
87
+ logger.info("[WS Thread] Listener loop finished.")
88
+ st.session_state['ws_connected'] = False
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89
 
90
  try:
91
+ asyncio.run(listen())
92
+ except Exception as e:
93
+ logger.error(f"[WS Thread] asyncio.run error: {e}")
94
+ logger.info("[WS Thread] Listener thread exiting.")
95
+
96
+
97
+ # --- Streamlit UI ---
98
+
99
+ st.set_page_config(layout="wide")
100
+
101
+ # --- Initialize Session State ---
102
+ if 'logged_in' not in st.session_state:
103
+ st.session_state.logged_in = False
104
+ if 'token' not in st.session_state:
105
+ st.session_state.token = None
106
+ if 'user_email' not in st.session_state:
107
+ st.session_state.user_email = None
108
+ if 'notifications' not in st.session_state:
109
+ st.session_state.notifications = []
110
+ if 'ws_thread' not in st.session_state:
111
+ st.session_state.ws_thread = None
112
+ if 'ws_connected' not in st.session_state:
113
+ st.session_state.ws_connected = False
114
+
115
+ # --- Notification Processing ---
116
+ new_notifications = []
117
+ while not notification_queue.empty():
118
+ try:
119
+ msg = notification_queue.get_nowait()
120
+ new_notifications.append(msg)
121
+ except queue.Empty:
122
+ break
123
+
124
+ if new_notifications:
125
+ logger.info(f"Processing {len(new_notifications)} notifications from queue.")
126
+ # Prepend new notifications to the session state list
127
+ current_list = st.session_state.notifications
128
+ st.session_state.notifications = new_notifications + current_list
129
+ # Limit history
130
+ if len(st.session_state.notifications) > 15:
131
+ st.session_state.notifications = st.session_state.notifications[:15]
132
+ # No explicit rerun needed here, Streamlit should rerun due to state change (?)
133
+ # or due to autorefresh below.
134
+
135
+ # --- Auto Refresh ---
136
+ # Refresh every 2 seconds to check the queue and update display
137
+ count = st_autorefresh(interval=2000, limit=None, key="notifrefresh")
138
+
139
+ # --- API Client ---
140
+ client = httpx.AsyncClient(base_url=API_BASE_URL, timeout=10.0)
141
+
142
+ # --- Helper Functions for API Calls ---
143
+ async def api_register(email, password):
144
+ try:
145
+ response = await client.post("/register", json={"email": email, "password": password})
146
+ response.raise_for_status()
147
+ return {"success": True, "data": response.json()}
148
+ except httpx.HTTPStatusError as e:
149
+ detail = e.response.json().get("detail", e.response.text)
150
+ logger.error(f"API Register Error: {e.response.status_code} - {detail}")
151
+ return {"success": False, "error": f"API Error: {detail}"}
152
+ except Exception as e:
153
+ logger.exception("Register call failed")
154
+ return {"success": False, "error": f"Request failed: {e}"}
155
+
156
+ async def api_login(email, password):
157
+ try:
158
+ response = await client.post("/login", json={"email": email, "password": password})
159
+ response.raise_for_status()
160
+ return {"success": True, "data": response.json()}
161
+ except httpx.HTTPStatusError as e:
162
+ detail = e.response.json().get("detail", e.response.text)
163
+ logger.error(f"API Login Error: {e.response.status_code} - {detail}")
164
+ return {"success": False, "error": f"API Error: {detail}"}
165
+ except Exception as e:
166
+ logger.exception("Login call failed")
167
+ return {"success": False, "error": f"Request failed: {e}"}
168
+
169
+ # --- UI Rendering ---
170
+ st.title("Authentication & Notification App (Streamlit)")
171
+
172
+ if not st.session_state.logged_in:
173
+ st.sidebar.header("Login or Register")
174
+ login_tab, register_tab = st.sidebar.tabs(["Login", "Register"])
175
+
176
+ with login_tab:
177
+ with st.form("login_form"):
178
+ login_email = st.text_input("Email", key="login_email")
179
+ login_password = st.text_input("Password", type="password", key="login_password")
180
+ login_button = st.form_submit_button("Login")
181
+
182
+ if login_button:
183
+ if not login_email or not login_password:
184
+ st.error("Please enter email and password.")
185
+ else:
186
+ result = asyncio.run(api_login(login_email, login_password)) # Run async in sync context
187
+ if result["success"]:
188
+ token = result["data"]["access_token"]
189
+ # Attempt to get user info immediately - needs modification if /users/me requires auth header
190
+ # For simplicity, just store email from login form for now
191
+ st.session_state.logged_in = True
192
+ st.session_state.token = token
193
+ st.session_state.user_email = login_email # Store email used for login
194
+ st.session_state.notifications = [] # Clear old notifications
195
+
196
+ # Start WebSocket listener thread
197
+ stop_event.clear() # Ensure stop event is clear
198
+ thread = threading.Thread(target=websocket_listener, args=(token,), daemon=True)
199
+ st.session_state.ws_thread = thread
200
+ thread.start()
201
+ logger.info("Login successful, WS thread started.")
202
+ st.rerun() # Rerun immediately to switch view
203
+ else:
204
+ st.error(f"Login failed: {result['error']}")
205
+
206
+ with register_tab:
207
+ with st.form("register_form"):
208
+ reg_email = st.text_input("Email", key="reg_email")
209
+ reg_password = st.text_input("Password", type="password", key="reg_password")
210
+ reg_confirm = st.text_input("Confirm Password", type="password", key="reg_confirm")
211
+ register_button = st.form_submit_button("Register")
212
+
213
+ if register_button:
214
+ if not reg_email or not reg_password or not reg_confirm:
215
+ st.error("Please fill all fields.")
216
+ elif reg_password != reg_confirm:
217
+ st.error("Passwords do not match.")
218
+ elif len(reg_password) < 8:
219
+ st.error("Password must be at least 8 characters.")
220
+ else:
221
+ result = asyncio.run(api_register(reg_email, reg_password))
222
+ if result["success"]:
223
+ st.success(f"Registration successful for {result['data']['email']}! Please log in.")
224
+ else:
225
+ st.error(f"Registration failed: {result['error']}")
226
+
227
+ else: # Logged In View
228
+ st.sidebar.header(f"Welcome, {st.session_state.user_email}!")
229
+ if st.sidebar.button("Logout"):
230
+ logger.info("Logout requested.")
231
+ # Stop WebSocket thread
232
+ if st.session_state.ws_thread and st.session_state.ws_thread.is_alive():
233
+ logger.info("Signalling WS thread to stop.")
234
+ stop_event.set()
235
+ st.session_state.ws_thread.join(timeout=2.0) # Wait briefly for thread exit
236
+ if st.session_state.ws_thread.is_alive():
237
+ logger.warning("WS thread did not exit cleanly.")
238
+ # Clear session state
239
+ st.session_state.logged_in = False
240
+ st.session_state.token = None
241
+ st.session_state.user_email = None
242
+ st.session_state.notifications = []
243
+ st.session_state.ws_thread = None
244
+ st.session_state.ws_connected = False
245
+ logger.info("Session cleared.")
246
+ st.rerun()
247
+
248
+ st.header("Dashboard")
249
+ # Display notifications
250
+ st.subheader("Real-time Notifications")
251
+ ws_status = "Connected" if st.session_state.ws_connected else "Disconnected"
252
+ st.caption(f"WebSocket Status: {ws_status}")
253
+
254
+ if st.session_state.notifications:
255
+ for i, msg in enumerate(st.session_state.notifications):
256
+ st.info(f"{msg}", icon="🔔")
257
+ else:
258
+ st.text("No new notifications.")
259
+
260
+ # Add a button to manually check queue/refresh if needed
261
+ # if st.button("Check for notifications"):
262
+ # st.rerun() # Force rerun which includes queue check
263
+
264
+
265
+ # --- Final Cleanup ---
266
+ # Ensure httpx client is closed if script exits abnormally
267
+ # (This might not always run depending on how Streamlit terminates)
268
+ # Ideally handled within context managers if used more extensively
269
+ # asyncio.run(client.aclose())
 
 
 
 
 
 
 
requirements.txt CHANGED
@@ -1,12 +1,15 @@
1
  fastapi==0.111.0
2
  uvicorn[standard]==0.29.0
3
- gradio==4.29.0
 
4
  passlib[bcrypt]==1.7.4
5
- bcrypt==4.1.3 # Add this line (or latest compatible version)
6
  python-dotenv==1.0.1
7
- databases[sqlite]==0.9.0 # Async DB access
8
- sqlalchemy==2.0.29 # Core needed by `databases` or for potential future use
9
  pydantic==2.7.1
10
- python-multipart==0.0.9 # For form data in FastAPI
11
- itsdangerous==2.1.2 # Simple secure signing for session IDs
12
- websockets>=11.0.3,<13.0 # Ensure compatibility
 
 
 
1
  fastapi==0.111.0
2
  uvicorn[standard]==0.29.0
3
+ # gradio==4.29.0 # Removed
4
+ streamlit==1.33.0 # Added (or latest)
5
  passlib[bcrypt]==1.7.4
6
+ bcrypt==4.1.3
7
  python-dotenv==1.0.1
8
+ databases[sqlite]==0.9.0
9
+ sqlalchemy==2.0.29
10
  pydantic==2.7.1
11
+ python-multipart==0.0.9
12
+ itsdangerous==2.1.2
13
+ websockets>=11.0.3,<13.0
14
+ httpx==0.27.0 # Needed for API calls from Streamlit
15
+ streamlit-autorefresh==1.0.1 # Added for polling notifications