support for arize phoenix ui to debug llm testing in notebook

dbtunnel/__init__.py

@@ -85,5 +85,10 @@ def uvicorn(app, port: int = 8080):
         from dbtunnel.uvicorn import UvicornAppTunnel
         return UvicornAppTunnel(app, port)
 
+    @staticmethod
+    def arize_phoenix(port: int = 9098):
+        from dbtunnel.arize_phoenix_ui import ArizePhoenixUITunnel
+        return ArizePhoenixUITunnel(port)
+
 
 dbtunnel = AppTunnels()

dbtunnel/arize_phoenix_ui.py

@@ -0,0 +1,67 @@
+import os
+import subprocess
+import sys
+
+from dbtunnel.tunnels import DbTunnel, DbTunnelProxy
+from dbtunnel.utils import execute
+from dbtunnel.vendor.asgiproxy.frameworks import Frameworks
+
+
+class ArizePhoenixUITunnel(DbTunnel):
+
+    def __init__(self, port: int):
+        super().__init__(port, "arize-phoenix")
+
+
+    def _imports(self):
+        try:
+            import phoenix.server
+            import nest_asyncio
+        except ImportError as e:
+            self._log.info("ImportError: Make sure you have chainlit installed. \n"
+                           "pip install arize-phoenix nest_asyncio")
+            raise e
+
+    def _display_url(self):
+        # must end with a "/" for it to not redirect
+        return f'<a href="{self._proxy_settings.proxy_url}">Click to go to {self._flavor} App!</a>'
+
+    def _run(self):
+        self.display()
+        self._log.info("Starting server...")
+
+        phoenix_service_port_no_share = 9099
+        proxy_service = None
+        if self._share is False:
+            url_base_path = self._proxy_settings.url_base_path
+            port = self._port
+
+            proxy_service = DbTunnelProxy(
+                proxy_port=port,
+                service_port=phoenix_service_port_no_share,
+                url_base_path=url_base_path,
+                framework=Frameworks.ARIZE_PHOENIX,
+                token_auth=self._basic_tunnel_auth["token_auth"],
+                token_auth_workspace_url=self._basic_tunnel_auth["token_auth_workspace_url"],
+            )
+
+            proxy_service.start()
+
+        my_env = os.environ.copy()
+        subprocess.run(f"kill -9 $(lsof -t -i:{self._port})", capture_output=True, shell=True)
+
+        if self.shared is False:
+            self._log.info(f"Deploying stable diffusion web ui app at path: \n{self._proxy_settings.proxy_url}")
+
+        if self._share is False:
+            cmd = [sys.executable, "-m", "-f", "phoenix.server.main", "--port", f"{phoenix_service_port_no_share}",
+                   "serve"]
+        else:
+            cmd = [sys.executable, "-m", "-f", "phoenix.server.main", "--port", f"{self._port}", "serve"]
+
+        self._log.info(f"Running command: {' '.join(cmd)}")
+        for path in execute(cmd, my_env):
+            self._log.info(path)
+
+        if proxy_service is not None:
+            proxy_service.wait()

dbtunnel/tunnels.py

@@ -92,7 +92,7 @@ def get_cloud_proxy_settings(cloud: str, org_id: str, cluster_id: str, port: int
 
 Flavor = Literal[
     "gradio", "fastapi", "nicegui", "streamlit", "stable-diffusion-ui", "bokeh", "flask", "dash", "solara",
-    "code-server", "chainlit", "shiny-python", "uvicorn"]
+    "code-server", "chainlit", "shiny-python", "uvicorn", "arize-phoenix"]
 
 
 def get_current_username() -> str:

dbtunnel/vendor/asgiproxy/frameworks.py

@@ -133,14 +133,59 @@ def _modify_js_bundle(content, root_path):
     )()
     return config
 
+
+def _make_arize_phoenix_local_proxy_config(
+        url_base_path,
+        service_host: str = "0.0.0.0",
+        service_port: int = 9989,
+        auth_config: dict = None
+):
+    auth_config = auth_config or {}
+
+    def _modify_root(content, root_path):
+        list_of_uris = [b"/index.css", b"/modernizr.js", b"/favicon.ico", b"/index.js", b"/graphql", b"/projects"]
+        for uri in list_of_uris:
+            content = content.replace(uri, root_path.encode("utf-8") + uri)
+        return content
+
+    def _modify_js_bundle(content, root_path):
+        list_of_uris = [b"/graphql", b"/projects"]
+        for uri in list_of_uris:
+            content = content.replace(uri, root_path.rstrip("/").encode("utf-8") + uri)
+
+        return content
+
+    modify_root = functools.partial(_modify_root, root_path=url_base_path)
+    modify_js_bundle = functools.partial(_modify_js_bundle, root_path=url_base_path)
+
+    config = type(
+        "Config",
+        (BaseURLProxyConfigMixin, ProxyConfig),
+        {
+            "upstream_base_url": f"http://{service_host}:{service_port}",
+            "rewrite_host_header": f"{service_host}:{service_port}",
+            "modify_content": {
+                "/": modify_root,
+                "/projects/": modify_root,
+                "/projects/*": modify_root,
+                "*/index.js": modify_js_bundle,
+                # some reason gradio also has caps index bundled calling out explicitly
+            },
+            **auth_config,
+        },
+    )()
+    return config
+
 class Frameworks:
     STREAMLIT: str = "streamlit"
     GRADIO: str = "gradio"
     CHAINLIT: str = "chainlit"
+    ARIZE_PHOENIX: str = "arize_phoenix"
 
 
 framework_specific_proxy_config = {
     Frameworks.STREAMLIT: _make_streamlit_local_proxy_config,
     Frameworks.GRADIO: _make_gradio_local_proxy_config,
     Frameworks.CHAINLIT: _make_chainlit_local_proxy_config,
+    Frameworks.ARIZE_PHOENIX: _make_arize_phoenix_local_proxy_config,
 }

dbtunnel/vendor/asgiproxy/proxies/http.py

@@ -91,6 +91,9 @@ async def convert_proxy_response_to_user_response(
                     response_content = modify_func(response_content)
                     new_headers.popall("Content-Length", None)
                     new_headers["Content-Length"] = str(len(response_content))
+                    # TODO: this may cause bugs :\ if we need multiple passes
+                    # TODO: in future maybe we have priority or flag
+                    break
 
     return Response(
         content=response_content,

setup.py

@@ -89,6 +89,9 @@
         "cli": [
             "click",
         ],
+        "arize-phoenix": [
+            "arize-phoenix"
+        ]
     },
     classifiers=[
         "Programming Language :: Python :: 3",