Upload tunnel_manager.py with huggingface_hub

tunnel_manager.py

@@ -0,0 +1,118 @@
+from pyngrok import ngrok, conf
+import psutil
+import time
+import os
+import json
+import yaml
+import requests
+from requests.exceptions import ConnectionError
+
+class TunnelManager:
+    DEFAULT_PORT = 5000
+    _current_tunnel = None
+    _ngrok_process = None
+    
+    @staticmethod
+    def load_config():
+        """Carga la configuración desde config.yaml"""
+        try:
+            config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
+            with open(config_path, 'r') as f:
+                return yaml.safe_load(f)
+        except Exception as e:
+            print(f"Error al cargar config.yaml: {e}")
+            return {}
+
+    @staticmethod
+    def setup_ngrok():
+        """Configura ngrok con el token del archivo config.yaml"""
+        config = TunnelManager.load_config()
+        token = config.get('NGROK_TOKEN')
+        
+        if token:
+            try:
+                ngrok.set_auth_token(token)
+                return True
+            except Exception as e:
+                print(f"Error al configurar el token de ngrok: {e}")
+        return False
+
+    @staticmethod
+    def cleanup():
+        """Limpia todas las sesiones de ngrok existentes"""
+        try:
+            print("\nLimpiando sesiones anteriores...")
+            # Intentar matar procesos de ngrok usando psutil
+            for proc in psutil.process_iter(['pid', 'name']):
+                try:
+                    if 'ngrok' in proc.name().lower():
+                        proc.kill()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+            
+            # Usar el método de pyngrok para limpiar
+            ngrok.kill()
+            time.sleep(2)  # Dar tiempo para que los procesos terminen
+            print("✓ Limpieza completada")
+            TunnelManager._current_tunnel = None
+        except Exception as e:
+            print(f"Error durante la limpieza: {e}")
+
+    @classmethod
+    def get_active_tunnel(cls):
+        """Obtiene la URL del túnel activo más reciente"""
+        try:
+            # Intentar diferentes puertos comunes de ngrok
+            ports = [4040, 4041, 5000]
+            for port in ports:
+                try:
+                    response = requests.get(f"http://127.0.0.1:{port}/api/tunnels", timeout=1)
+                    if response.status_code == 200:
+                        tunnels = response.json()["tunnels"]
+                        if tunnels:
+                            for tunnel in tunnels:
+                                if tunnel["proto"] == "https":
+                                    cls._current_tunnel = tunnel["public_url"]
+                                    return cls._current_tunnel
+                except ConnectionError:
+                    continue
+                except Exception:
+                    continue
+            return None
+        except Exception as e:
+            print(f"No se encontraron túneles activos")
+            return None
+
+    @classmethod
+    def start(cls, port=5000):
+        """Inicia un túnel de ngrok"""
+        if cls._current_tunnel:
+            return cls._current_tunnel
+            
+        try:
+            # Asegurarse de que no hay túneles activos
+            cls.cleanup()
+            
+            # Configurar e iniciar nuevo túnel
+            tunnel = ngrok.connect(port, "http")
+            cls._current_tunnel = tunnel.public_url
+            
+            # Verificar que el túnel está activo
+            time.sleep(1)
+            tunnels = ngrok.get_tunnels()
+            if not tunnels:
+                raise Exception("No se pudo establecer el túnel")
+                
+            return cls._current_tunnel
+            
+        except Exception as e:
+            print(f"Error iniciando túnel: {e}")
+            return None
+
+    @staticmethod
+    def get_tunnels():
+        """Obtiene la lista de túneles activos"""
+        try:
+            return ngrok.get_tunnels()
+        except:
+            return []