From d055e1f88809145fe5c45c4f4faa811f09e8a665 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sun, 22 Sep 2024 02:14:59 +0200
Subject: [PATCH] refac

---
 backend/open_webui/apps/socket/main.py | 14 +++++++++++++-
 backend/open_webui/env.py              |  1 +
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/backend/open_webui/apps/socket/main.py b/backend/open_webui/apps/socket/main.py
index c0577e648..9a3db51e5 100644
--- a/backend/open_webui/apps/socket/main.py
+++ b/backend/open_webui/apps/socket/main.py
@@ -1,5 +1,7 @@
 import asyncio
 import socketio
+import logging
+import sys
 import time
 
 from open_webui.apps.webui.models.users import Users
@@ -11,6 +13,16 @@ from open_webui.env import (
 from open_webui.utils.utils import decode_token
 from open_webui.apps.socket.utils import RedisDict
 
+from open_webui.env import (
+    GLOBAL_LOG_LEVEL,
+    SRC_LOG_LEVELS,
+)
+
+
+logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
+log = logging.getLogger(__name__)
+log.setLevel(SRC_LOG_LEVELS["SOCKET"])
+
 
 if WEBSOCKET_MANAGER == "redis":
     mgr = socketio.AsyncRedisManager(WEBSOCKET_REDIS_URL)
@@ -57,7 +69,7 @@ TIMEOUT_DURATION = 3
 async def periodic_usage_pool_cleanup():
     while True:
         now = int(time.time())
-        print("Cleaning up usage pool", now)
+        log.debug("Cleaning up usage pool", now)
         for model_id, connections in list(USAGE_POOL.items()):
             # Creating a list of sids to remove if they have timed out
             expired_sids = [
diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py
index 504eeea54..4536f91fd 100644
--- a/backend/open_webui/env.py
+++ b/backend/open_webui/env.py
@@ -84,6 +84,7 @@ log_sources = [
     "OPENAI",
     "RAG",
     "WEBHOOK",
+    "SOCKET",
 ]
 
 SRC_LOG_LEVELS = {}