initial implementation

Author: codelion

optillm/plugins/proxy/README.md

@@ -0,0 +1,8 @@
+"""OptiLLM Proxy Plugin - Load balancing and failover for LLM providers"""
+
+from .config import ProxyConfig
+from .client import ProxyClient
+from .routing import RouterFactory
+from .health import HealthChecker
+
+__all__ = ['ProxyConfig', 'ProxyClient', 'RouterFactory', 'HealthChecker']

optillm/plugins/proxy/__init__.py

@@ -0,0 +1,182 @@
+"""
+Dynamic handler for approaches and plugins - no hardcoding.
+"""
+import importlib
+import importlib.util
+import logging
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+class ApproachHandler:
+    """Dynamically handles both approaches and plugins"""
+    
+    def __init__(self):
+        self._approaches_cache = {}
+        self._plugins_cache = {}
+        self._discovered = False
+    
+    def handle(self, name: str, system_prompt: str, initial_query: str, 
+               client, model: str, request_config: dict = None) -> Optional[Tuple[str, int]]:
+        """
+        Try to handle the given name as an approach or plugin.
+        Returns None if not found, otherwise returns (response, tokens)
+        """
+        # Lazy discovery
+        if not self._discovered:
+            self._discover_handlers()
+            self._discovered = True
+        
+        # Check if it's an approach
+        if name in self._approaches_cache:
+            logger.info(f"Routing approach '{name}' through proxy")
+            handler = self._approaches_cache[name]
+            return self._execute_handler(
+                handler, system_prompt, initial_query, client, model, request_config
+            )
+        
+        # Check if it's a plugin
+        if name in self._plugins_cache:
+            logger.info(f"Routing plugin '{name}' through proxy")
+            handler = self._plugins_cache[name]
+            return self._execute_handler(
+                handler, system_prompt, initial_query, client, model, request_config
+            )
+        
+        logger.debug(f"'{name}' not recognized as approach or plugin")
+        return None
+    
+    def _discover_handlers(self):
+        """Discover available approaches and plugins dynamically"""
+        
+        # Discover approaches
+        self._discover_approaches()
+        
+        # Discover plugins
+        self._discover_plugins()
+        
+        logger.info(f"Discovered {len(self._approaches_cache)} approaches, "
+                   f"{len(self._plugins_cache)} plugins")
+    
+    def _discover_approaches(self):
+        """Discover built-in approaches from optillm package"""
+        approach_modules = {
+            'mcts': ('optillm.mcts', 'chat_with_mcts'),
+            'bon': ('optillm.bon', 'best_of_n_sampling'),
+            'moa': ('optillm.moa', 'mixture_of_agents'),
+            'rto': ('optillm.rto', 'round_trip_optimization'),
+            'self_consistency': ('optillm.self_consistency', 'advanced_self_consistency_approach'),
+            'pvg': ('optillm.pvg', 'inference_time_pv_game'),
+            'z3': ('optillm.z3_solver', None),  # Special case
+            'rstar': ('optillm.rstar', None),  # Special case
+            'cot_reflection': ('optillm.cot_reflection', 'cot_reflection'),
+            'plansearch': ('optillm.plansearch', 'plansearch'),
+            'leap': ('optillm.leap', 'leap'),
+            're2': ('optillm.reread', 're2_approach'),
+            'cepo': ('optillm.cepo.cepo', 'cepo'),  # CEPO approach
+        }
+        
+        for name, (module_path, func_name) in approach_modules.items():
+            try:
+                module = importlib.import_module(module_path)
+                
+                if name == 'z3':
+                    # Special handling for Z3
+                    solver_class = getattr(module, 'Z3SymPySolverSystem')
+                    self._approaches_cache[name] = lambda s, q, c, m, **kw: \
+                        solver_class(s, c, m).process_query(q)
+                elif name == 'rstar':
+                    # Special handling for RStar
+                    rstar_class = getattr(module, 'RStar')
+                    self._approaches_cache[name] = lambda s, q, c, m, **kw: \
+                        rstar_class(s, c, m, **kw).solve(q)
+                elif name == 'cepo':
+                    # Special handling for CEPO which needs special config
+                    cepo_func = getattr(module, func_name)
+                    # We'll pass empty CepoConfig for now - it can be enhanced later
+                    self._approaches_cache[name] = cepo_func
+                else:
+                    if func_name:
+                        self._approaches_cache[name] = getattr(module, func_name)
+                    
+            except (ImportError, AttributeError) as e:
+                logger.debug(f"Could not load approach '{name}': {e}")
+    
+    def _discover_plugins(self):
+        """Discover available plugins dynamically"""
+        try:
+            import optillm
+            import os
+            import glob
+            
+            # Get plugin directories
+            package_dir = Path(optillm.__file__).parent / 'plugins'
+            
+            # Find all Python files in plugins directory
+            plugin_files = []
+            if package_dir.exists():
+                plugin_files.extend(glob.glob(str(package_dir / '*.py')))
+            
+            for plugin_file in plugin_files:
+                if '__pycache__' in plugin_file or '__init__' in plugin_file:
+                    continue
+                    
+                try:
+                    # Extract module name
+                    module_name = Path(plugin_file).stem
+                    
+                    # Skip self
+                    if module_name == 'proxy_plugin':
+                        continue
+                    
+                    # Import module
+                    spec = importlib.util.spec_from_file_location(module_name, plugin_file)
+                    if spec and spec.loader:
+                        module = importlib.util.module_from_spec(spec)
+                        spec.loader.exec_module(module)
+                        
+                        # Check if it has required attributes
+                        if hasattr(module, 'SLUG') and hasattr(module, 'run'):
+                            slug = getattr(module, 'SLUG')
+                            run_func = getattr(module, 'run')
+                            self._plugins_cache[slug] = run_func
+                            
+                except Exception as e:
+                    logger.debug(f"Could not load plugin from {plugin_file}: {e}")
+                    
+        except Exception as e:
+            logger.debug(f"Error discovering plugins: {e}")
+    
+    def _execute_handler(self, handler, system_prompt: str, initial_query: str,
+                        client, model: str, request_config: dict = None) -> Tuple[str, int]:
+        """Execute a handler function with proper signature detection"""
+        try:
+            # Check function signature
+            sig = inspect.signature(handler)
+            params = sig.parameters
+            
+            # Build arguments based on signature
+            args = [system_prompt, initial_query, client, model]
+            kwargs = {}
+            
+            # Check if handler accepts request_config
+            if 'request_config' in params:
+                kwargs['request_config'] = request_config
+            
+            # Some handlers may accept additional kwargs
+            if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()):
+                # Only add safe kwargs that won't conflict
+                if request_config:
+                    # Filter out parameters that might conflict
+                    safe_kwargs = {k: v for k, v in request_config.items() 
+                                 if k not in ['model', 'messages', 'system_prompt', 'initial_query']}
+                    kwargs.update(safe_kwargs)
+            
+            # Execute handler
+            return handler(*args, **kwargs)
+            
+        except Exception as e:
+            logger.error(f"Error executing handler: {e}")
+            raise

optillm/plugins/proxy/approach_handler.py

@@ -0,0 +1,186 @@
+"""
+ProxyClient implementation for load balancing across multiple LLM providers.
+"""
+import time
+import logging
+import random
+from typing import Dict, List, Any, Optional
+from openai import OpenAI, AzureOpenAI
+from .routing import RouterFactory
+from .health import HealthChecker
+
+logger = logging.getLogger(__name__)
+
+class Provider:
+    """Wrapper for a provider configuration and client"""
+    def __init__(self, config: Dict):
+        self.name = config['name']
+        self.base_url = config['base_url']
+        self.api_key = config['api_key']
+        self.weight = config.get('weight', 1)
+        self.fallback_only = config.get('fallback_only', False)
+        self.model_map = config.get('model_map', {})
+        self._client = None
+        self.is_healthy = True
+        self.last_error = None
+        self.latencies = []  # Track recent latencies
+        
+    @property
+    def client(self):
+        """Lazy initialization of OpenAI client"""
+        if not self._client:
+            if 'azure' in self.base_url.lower():
+                # Handle Azure OpenAI
+                self._client = AzureOpenAI(
+                    api_key=self.api_key,
+                    azure_endpoint=self.base_url,
+                    api_version="2024-02-01"
+                )
+            else:
+                # Standard OpenAI-compatible client
+                self._client = OpenAI(
+                    api_key=self.api_key,
+                    base_url=self.base_url
+                )
+        return self._client
+    
+    def map_model(self, model: str) -> str:
+        """Map requested model to provider-specific name"""
+        return self.model_map.get(model, model)
+    
+    def track_latency(self, latency: float):
+        """Track request latency"""
+        self.latencies.append(latency)
+        if len(self.latencies) > 10:
+            self.latencies.pop(0)
+    
+    def avg_latency(self) -> float:
+        """Get average latency"""
+        if not self.latencies:
+            return 0
+        return sum(self.latencies) / len(self.latencies)
+
+class ProxyClient:
+    """OpenAI-compatible client that proxies to multiple providers"""
+    
+    def __init__(self, config: Dict, fallback_client=None):
+        self.config = config
+        self.fallback_client = fallback_client
+        
+        # Initialize providers
+        self.providers = [
+            Provider(p) for p in config.get('providers', [])
+        ]
+        
+        # Filter out fallback-only providers for normal routing
+        self.active_providers = [
+            p for p in self.providers if not p.fallback_only
+        ]
+        
+        self.fallback_providers = [
+            p for p in self.providers if p.fallback_only
+        ]
+        
+        # Initialize router
+        strategy = config.get('routing', {}).get('strategy', 'round_robin')
+        self.router = RouterFactory.create(strategy, self.active_providers)
+        
+        # Initialize health checker
+        health_config = config.get('routing', {}).get('health_check', {})
+        self.health_checker = HealthChecker(
+            providers=self.providers,
+            enabled=health_config.get('enabled', True),
+            interval=health_config.get('interval', 30),
+            timeout=health_config.get('timeout', 5)
+        )
+        
+        # Start health checking
+        self.health_checker.start()
+        
+        # Monitoring settings
+        monitoring = config.get('monitoring', {})
+        self.track_latency = monitoring.get('track_latency', True)
+        self.track_errors = monitoring.get('track_errors', True)
+        
+        # Create chat namespace
+        self.chat = self._Chat(self)
+    
+    class _Chat:
+        def __init__(self, proxy_client):
+            self.proxy_client = proxy_client
+            self.completions = proxy_client._Completions(proxy_client)
+    
+    class _Completions:
+        def __init__(self, proxy_client):
+            self.proxy_client = proxy_client
+        
+        def create(self, **kwargs):
+            """Create completion with load balancing and failover"""
+            model = kwargs.get('model', 'unknown')
+            attempted_providers = set()
+            errors = []
+            
+            # Get healthy providers
+            healthy_providers = [
+                p for p in self.proxy_client.active_providers 
+                if p.is_healthy
+            ]
+            
+            if not healthy_providers:
+                logger.warning("No healthy providers, trying fallback providers")
+                healthy_providers = self.proxy_client.fallback_providers
+            
+            # Try routing through healthy providers
+            while healthy_providers:
+                available_providers = [p for p in healthy_providers if p not in attempted_providers]
+                if not available_providers:
+                    break
+                    
+                provider = self.proxy_client.router.select(available_providers)
+                
+                if not provider:
+                    break
+                    
+                attempted_providers.add(provider)
+                
+                try:
+                    # Map model name if needed
+                    request_kwargs = kwargs.copy()
+                    request_kwargs['model'] = provider.map_model(model)
+                    
+                    # Track timing
+                    start_time = time.time()
+                    
+                    # Make request
+                    logger.debug(f"Routing to {provider.name}")
+                    response = provider.client.chat.completions.create(**request_kwargs)
+                    
+                    # Track success
+                    latency = time.time() - start_time
+                    if self.proxy_client.track_latency:
+                        provider.track_latency(latency)
+                    
+                    logger.info(f"Request succeeded via {provider.name} in {latency:.2f}s")
+                    return response
+                    
+                except Exception as e:
+                    logger.error(f"Provider {provider.name} failed: {e}")
+                    errors.append((provider.name, str(e)))
+                    
+                    # Mark provider as unhealthy
+                    if self.proxy_client.track_errors:
+                        provider.is_healthy = False
+                        provider.last_error = str(e)
+            
+            # All providers failed, try fallback client
+            if self.proxy_client.fallback_client:
+                logger.warning("All proxy providers failed, using fallback client")
+                try:
+                    return self.proxy_client.fallback_client.chat.completions.create(**kwargs)
+                except Exception as e:
+                    errors.append(("fallback_client", str(e)))
+            
+            # Complete failure
+            error_msg = f"All providers failed. Errors: {errors}"
+            logger.error(error_msg)
+            raise Exception(error_msg)