@@ -142,6 +142,13 @@ async def admin_logout(request: Request):
142142_log_counter = itertools .count ()
143143_log_buffer : deque [Dict [str , Any ]] = deque (maxlen = 500 )
144144_ws_clients : List [WebSocket ] = []
145+ _event_loop : Optional [asyncio .AbstractEventLoop ] = None
146+
147+
148+ def _set_event_loop (loop : asyncio .AbstractEventLoop ) -> None :
149+ """Store the main event loop reference for cross-thread access."""
150+ global _event_loop
151+ _event_loop = loop
145152
146153
147154class WebSocketLogHandler (logging .Handler ):
@@ -160,13 +167,17 @@ def emit(self, record: logging.LogRecord) -> None:
160167
161168 _log_buffer .append (entry )
162169
163- # Broadcast to all connected WebSocket clients
164- # Use call_soon_threadsafe to safely schedule from any thread
170+ loop = _event_loop
171+ if loop is None or loop .is_closed ():
172+ return
173+
165174 for ws in list (_ws_clients ):
166175 try :
167- loop = asyncio .get_event_loop ()
168- loop .call_soon_threadsafe (lambda w = ws , e = entry : asyncio .create_task (_send_log_safe (w , e )))
169- except Exception :
176+ loop .call_soon_threadsafe (
177+ loop .create_task ,
178+ _send_log_safe (ws , entry ),
179+ )
180+ except RuntimeError :
170181 pass
171182
172183
@@ -175,19 +186,16 @@ async def _send_log_safe(websocket: WebSocket, entry: Dict[str, Any]) -> None:
175186 try :
176187 await websocket .send_json (entry )
177188 except Exception :
178- # Client disconnected or other error — ignore
179189 pass
180190
181191
182192# Install the handler on the root logger so ALL logs are captured
183193_ws_log_handler = WebSocketLogHandler ()
184194_ws_log_handler .setLevel (logging .INFO )
185195logging .getLogger ().addHandler (_ws_log_handler )
186- # Also capture xmem-specific loggers
187196logging .getLogger ("xmem" ).addHandler (_ws_log_handler )
188197logging .getLogger ("src" ).addHandler (_ws_log_handler )
189198logging .getLogger ("uvicorn" ).addHandler (_ws_log_handler )
190- # Explicitly capture AWS boto3 logs at INFO level
191199logging .getLogger ("boto3" ).setLevel (logging .INFO )
192200logging .getLogger ("botocore" ).setLevel (logging .INFO )
193201logging .getLogger ("boto3" ).addHandler (_ws_log_handler )
@@ -199,6 +207,9 @@ async def ws_live_logs(websocket: WebSocket):
199207 """WebSocket endpoint for live log streaming."""
200208 await websocket .accept ()
201209
210+ # Capture the running event loop so the log handler can broadcast
211+ _set_event_loop (asyncio .get_running_loop ())
212+
202213 # Validate auth token from query param
203214 token = websocket .query_params .get ("token" , "" )
204215 if token not in _admin_sessions :
@@ -208,27 +219,22 @@ async def ws_live_logs(websocket: WebSocket):
208219 _ws_clients .append (websocket )
209220
210221 try :
211- last_id = - 1
212222 # Send buffered logs first
213223 for entry in list (_log_buffer ):
214224 await websocket .send_json (entry )
215- last_id = entry .get ("id" , - 1 )
216225
217- # Keep alive — wait for disconnect or new logs
226+ # Keep alive — the WebSocketLogHandler.emit() pushes new logs
227+ # via call_soon_threadsafe. We just need to keep the connection
228+ # open by waiting for client messages (or disconnect).
218229 while True :
219- # Send any new logs that arrived in the buffer
220- for entry in list (_log_buffer ):
221- entry_id = entry .get ("id" , - 1 )
222- if entry_id > last_id :
223- await websocket .send_json (entry )
224- last_id = entry_id
225-
226230 try :
227- # Sleep briefly, check if client disconnected
228- await asyncio .wait_for (websocket .receive_text (), timeout = 0.5 )
231+ await asyncio .wait_for (websocket .receive_text (), timeout = 30 )
229232 except asyncio .TimeoutError :
230- # Send ping to keep connection alive occasionally
231- pass
233+ # Send a lightweight ping to detect broken connections
234+ try :
235+ await websocket .send_json ({"type" : "ping" })
236+ except Exception :
237+ break
232238 except WebSocketDisconnect :
233239 pass
234240 except Exception :
@@ -238,6 +244,33 @@ async def ws_live_logs(websocket: WebSocket):
238244 _ws_clients .remove (websocket )
239245
240246
247+ # ═══════════════════════════════════════════════════════════════════════════
248+ # Gemini pricing (paid tier, per 1M tokens in USD)
249+ # ═══════════════════════════════════════════════════════════════════════════
250+ COST_TABLE : Dict [str , Dict [str , float ]] = {
251+ "gemini-embedding-001" : {"input" : 0.15 , "output" : 0.00 },
252+ "gemini-2.5-flash-lite" : {"input" : 0.10 , "output" : 0.40 },
253+ "gemini-2.5-flash" : {"input" : 0.15 , "output" : 0.60 },
254+ "gemini-2.0-flash" : {"input" : 0.10 , "output" : 0.40 },
255+ "gemini-2.0-flash-lite" : {"input" : 0.075 , "output" : 0.30 },
256+ }
257+ _PER_M = 1_000_000
258+
259+
260+ def _estimate_cost (model : str , input_tokens : int , output_tokens : int ) -> float :
261+ """Estimate USD cost for a given model and token counts."""
262+ key = model .lower ().strip ()
263+ prices = COST_TABLE .get (key )
264+ if prices is None :
265+ for prefix , p in COST_TABLE .items ():
266+ if key .startswith (prefix ):
267+ prices = p
268+ break
269+ if prices is None :
270+ return 0.0
271+ return (input_tokens * prices ["input" ] + output_tokens * prices ["output" ]) / _PER_M
272+
273+
241274# ═══════════════════════════════════════════════════════════════════════════
242275# Analytics summary API
243276# ═══════════════════════════════════════════════════════════════════════════
@@ -284,6 +317,15 @@ async def analytics_summary(request: Request, user: dict = Depends(_verify_admin
284317 {"$sort" : {"count" : - 1 }},
285318 ]))
286319
320+ # Compute cost for each LLM stats row
321+ for row in llm_stats_24h :
322+ model_name = (row .get ("_id" ) or {}).get ("model" , "" )
323+ row ["cost_usd" ] = round (_estimate_cost (
324+ model_name ,
325+ row .get ("total_input_tokens" , 0 ),
326+ row .get ("total_output_tokens" , 0 ),
327+ ), 6 )
328+
287329 # Hourly request volume (last 24h)
288330 hourly_volume = list (collection .aggregate ([
289331 {"$match" : {"event" : "api_call" , "ts" : {"$gte" : last_24h }}},
@@ -303,7 +345,7 @@ async def analytics_summary(request: Request, user: dict = Depends(_verify_admin
303345 "event" : "api_call" , "ts" : {"$gte" : last_24h }, "user_id" : {"$ne" : "" },
304346 })
305347
306- # Total token usage (last 7d, last 30d )
348+ # Total token usage (last 7d)
307349 token_usage_7d = list (collection .aggregate ([
308350 {"$match" : {"event" : "llm_call" , "ts" : {"$gte" : last_7d }}},
309351 {"$group" : {
@@ -315,6 +357,28 @@ async def analytics_summary(request: Request, user: dict = Depends(_verify_admin
315357 }},
316358 ]))
317359
360+ # Per-model cost breakdown (last 7d)
361+ cost_by_model_7d = list (collection .aggregate ([
362+ {"$match" : {"event" : "llm_call" , "ts" : {"$gte" : last_7d }}},
363+ {"$group" : {
364+ "_id" : "$model" ,
365+ "total_input" : {"$sum" : "$input_tokens" },
366+ "total_output" : {"$sum" : "$output_tokens" },
367+ "call_count" : {"$sum" : 1 },
368+ }},
369+ {"$sort" : {"total_input" : - 1 }},
370+ ]))
371+
372+ total_cost_7d = 0.0
373+ for row in cost_by_model_7d :
374+ model_name = row .get ("_id" ) or ""
375+ cost = _estimate_cost (model_name , row .get ("total_input" , 0 ), row .get ("total_output" , 0 ))
376+ row ["cost_usd" ] = round (cost , 6 )
377+ total_cost_7d += cost
378+
379+ if token_usage_7d :
380+ token_usage_7d [0 ]["total_cost_usd" ] = round (total_cost_7d , 6 )
381+
318382 # Daily LLM calls (last 7d) for chart
319383 daily_llm = list (collection .aggregate ([
320384 {"$match" : {"event" : "llm_call" , "ts" : {"$gte" : last_7d }}},
@@ -337,6 +401,8 @@ async def analytics_summary(request: Request, user: dict = Depends(_verify_admin
337401 "unique_users_24h" : len (unique_users ),
338402 "token_usage_7d" : _bson_safe (token_usage_7d [0 ] if token_usage_7d else {}),
339403 "daily_llm_calls" : _bson_safe (daily_llm ),
404+ "cost_by_model_7d" : _bson_safe (cost_by_model_7d ),
405+ "cost_table" : COST_TABLE ,
340406 })
341407 except Exception as exc :
342408 logger .exception ("Analytics summary failed" )
0 commit comments