feat: Separate webui from openai-forward main program

.env

@@ -18,9 +18,9 @@ DEFAULT_REQUEST_CACHING_VALUE=false
 
 #BENCHMARK_MODE=true
 
-FORWARD_CONFIG=[{"base_url":"https://api.openai.com","route":"/","type":"openai"}]
+FORWARD_CONFIG=[{"base_url":"https://api.openai-forward.com","route":"/","type":"openai"}]
 
-#OPENAI_API_KEY=
+#OPENAI_API_KEY={"sk-xxx": 0, "sk-xxx": 1, "sk-xxx": 2}
 #FORWARD_KEY={"fk-0": 0, "fk-1": 1, "fk-2": 2}
 #LEVEL_MODELS={"1": ["gpt-4"], "2": ["gpt-3.5-turbo"]}
 

docker-compose.yaml

@@ -1,6 +1,6 @@
 services:
   openai_forward:
-    image: beidongjiedeguang/openai-forward:latest
+    image: beidongjiedeguang/openai-forward:webui-latest
     container_name: openai-forward-container
     env_file:
       .env

openai_forward/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.7.1"
+__version__ = "0.7.2"
 
 from dotenv import load_dotenv
 

openai_forward/__main__.py

@@ -11,15 +11,35 @@
 
 
 class Cli:
-    def run(self, port=8000, workers=1, webui=False, ui_port=8001):
+    def run_web(self, port=8001, openai_forward_host='localhost', wait=True):
+        """
+        Runs the web UI using the Streamlit server.
+
+        Args:
+            port (int): The port number on which to run the server.
+            openai_forward_host (str): The host of the OpenAI Forward server.
+            wait (bool): Whether to wait for the server to stop. Default is True.
+
+        Returns:
+            None
+        """
+        os.environ['OPENAI_FORWARD_HOST'] = openai_forward_host
+        try:
+            self._start_streamlit(port=port, wait=wait)
+        except KeyboardInterrupt:
+            ...
+        except Exception as e:
+            raise
+
+    def run(self, port=8000, workers=1, webui=False, start_ui=True, ui_port=8001):
         """
         Runs the application using the Uvicorn server.
 
         Args:
-            port (int): The port number on which to run the server. Default is 8000.
-            workers (int): The number of worker processes to run. Default is 1.
+            port (int): The port number on which to run the server.
+            workers (int): The number of worker processes to run.
             webui (bool): Whether to run the web UI. Default is False.
-            ui_port (int): The port number on which to run streamlit. Default is 17860.
+            ui_port (int): The port number on which to run streamlit.
 
         Returns:
             None
@@ -42,28 +62,54 @@ def run(self, port=8000, workers=1, webui=False, ui_port=8001):
                 ssl_certfile=ssl_certfile,
             )
         else:
-            os.environ['OPENAI_FORWARD_WEBUI'] = 'true'
+            import threading
 
             import zmq
+            from flaxkv.helper import SimpleQueue
+
+            from openai_forward.helper import get_inner_ip
 
             mq_port = 15555
+
+            os.environ['OPENAI_FORWARD_WEBUI'] = 'true'
+
             context = zmq.Context()
             socket = context.socket(zmq.REP)
             socket.bind(f"tcp://*:{mq_port}")
+            log_socket = context.socket(zmq.ROUTER)
+            log_socket.bind(f"tcp://*:{15556}")
+            subscriber_info = {}
+
+            def mq_worker(log_socket: zmq.Socket):
+
+                while True:
+                    identity, uid, message = log_socket.recv_multipart()
+                    if uid == b"/subscribe":
+                        subscriber_info[identity] = True
+                        continue
+                    else:
+                        for subscriber, _ in subscriber_info.items():
+                            log_socket.send_multipart([subscriber, uid, message])
+
+            thread = threading.Thread(target=mq_worker, args=(log_socket,))
+            thread.daemon = True
+            thread.start()
 
             self._start_uvicorn(
                 port=port,
                 workers=workers,
                 ssl_keyfile=ssl_keyfile,
                 ssl_certfile=ssl_certfile,
             )
-            self._start_streamlit(port=ui_port)
-            atexit.register(self._stop)
+
+            if start_ui:
+                self._start_streamlit(port=ui_port, wait=False)
+
+            atexit.register(self._stop_uvicorn)
 
             while True:
                 message = socket.recv()
                 env_dict: dict = pickle.loads(message)
-                # logger.debug(f"{env_dict=}")
 
                 for key, value in env_dict.items():
                     os.environ[key] = value
@@ -102,7 +148,7 @@ def _start_uvicorn(self, port, workers, ssl_keyfile=None, ssl_certfile=None):
             suppress_exception=suppress_exception,
         )
 
-    def _start_streamlit(self, port):
+    def _start_streamlit(self, port, wait=False):
         from openai_forward.helper import relp
 
         self.streamlit_proc = subprocess.Popen(
@@ -125,10 +171,20 @@ def _start_streamlit(self, port):
             ]
         )
 
+        atexit.register(self._stop_streamlit)
+        if wait:
+            self.streamlit_proc.wait()
+
     def _restart_uvicorn(self, **kwargs):
-        self._stop(streamlit=False)
+        self._stop_uvicorn()
         self._start_uvicorn(**kwargs)
 
+    def _stop_streamlit(self):
+        self._stop(uvicorn=False)
+
+    def _stop_uvicorn(self):
+        self._stop(streamlit=False)
+
     def _stop(self, uvicorn=True, streamlit=True):
         if uvicorn and self.uvicorn_proc.poll() is None:
             self.uvicorn_proc.send_signal(signal.SIGINT)

openai_forward/content/openai.py

@@ -21,15 +21,15 @@ def __init__(self, route_prefix: str, _suffix: str):
         self.logger = logger.bind(**kwargs)
 
         self.webui = False
-        if os.environ.get("OPENAI_FORWARD_WEBUI"):
+        if os.environ.get("OPENAI_FORWARD_WEBUI").strip().lower() == 'true':
             self.webui = True
 
             import zmq
             from flaxkv.helper import SimpleQueue
 
             context = zmq.Context()
             socket = context.socket(zmq.DEALER)
-            socket.connect(f"tcp://localhost:15556")
+            socket.connect("tcp://localhost:15556")
 
             self.q = SimpleQueue(maxsize=200)
 

openai_forward/helper.py

@@ -3,6 +3,7 @@
 import inspect
 import os
 import re
+import socket
 import time
 from pathlib import Path
 from typing import Dict, List, Union
@@ -17,6 +18,12 @@ def __contains__(self, item):
         return True
 
 
+def get_inner_ip():
+    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+        s.connect(('8.8.8.8', 80))
+        return s.getsockname()[0]
+
+
 def urljoin(base_url, *relative_urls):
     """
     This function concatenates a base URL with any number of relative URL segments, producing a complete URL string.

openai_forward/webui/chat.py

@@ -10,7 +10,9 @@ def render_chat_log_message(msg: Dict):
         with st.chat_message(name="human"):
             messages = msg.pop('messages')
             for msg_item in messages:
-                st.write(f"`{msg_item['role']}`: {msg_item['content']}")
+                # https://github.com/streamlit/streamlit/issues/7978
+                # st.write(f"`{msg_item['role']}`: {msg_item['content']}")
+                st.text(f"`{msg_item['role']}`: {msg_item['content']}")
             st.write(msg)
     elif msg.get("assistant_role"):
         with st.chat_message(name="ai"):

openai_forward/webui/run.py

@@ -1,3 +1,4 @@
+import os
 import pickle
 import threading
 import time
@@ -30,23 +31,23 @@ def get_global_vars():
     context = zmq.Context()
     socket = context.socket(zmq.REQ)
     # socket.setsockopt(zmq.CONNECT_TIMEOUT, 20000) # 20s
-    log_socket = context.socket(zmq.ROUTER)
+    openai_forward_host = os.environ.get("OPENAI_FORWARD_HOST", "localhost")
+    socket.connect(f"tcp://{openai_forward_host}:15555")
 
-    socket.connect("tcp://localhost:15555")
-    log_socket.bind("tcp://*:15556")
+    log_socket = context.socket(zmq.DEALER)
+    log_socket.connect(f"tcp://{openai_forward_host}:15556")
+    log_socket.send_multipart([b"/subscribe", b"0"])
 
     def worker(log_socket: zmq.Socket, q: SimpleQueue):
         while True:
             message = log_socket.recv_multipart()
-            # print(f"{message=}")
-            identify, uid, msg = message
+            uid, msg = message
             q.put((uid, msg))
 
-    q = SimpleQueue(maxsize=200)
+    q = SimpleQueue(maxsize=100)
     threading.Thread(target=worker, args=(log_socket, q)).start()
     config = Config().come_from_env()
-    # print(f"{config=}")
-    chat_data = ChatData(200, render_chat_log_message)
+    chat_data = ChatData(100, render_chat_log_message)
     return {
         "socket": socket,
         "log_socket": log_socket,