fix proxy

Author: codelion

optillm/plugins/proxy/__init__.py

@@ -1,8 +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
+from optillm.plugins.proxy.config import ProxyConfig
+from optillm.plugins.proxy.client import ProxyClient
+from optillm.plugins.proxy.routing import RouterFactory
+from optillm.plugins.proxy.health import HealthChecker
 
 __all__ = ['ProxyConfig', 'ProxyClient', 'RouterFactory', 'HealthChecker']

optillm/plugins/proxy/client.py

@@ -6,8 +6,8 @@
 import random
 from typing import Dict, List, Any, Optional
 from openai import OpenAI, AzureOpenAI
-from .routing import RouterFactory
-from .health import HealthChecker
+from optillm.plugins.proxy.routing import RouterFactory
+from optillm.plugins.proxy.health import HealthChecker
 
 logger = logging.getLogger(__name__)
 
@@ -137,6 +137,7 @@ def create(self, **kwargs):
                     break
 
                 provider = self.proxy_client.router.select(available_providers)
+                logger.info(f"Router selected provider: {provider.name if provider else 'None'}")
 
                 if not provider:
                     break

optillm/plugins/proxy/routing.py

@@ -2,9 +2,16 @@
 Routing strategies for load balancing across providers.
 """
 import random
+import logging
 from typing import List, Optional
 from abc import ABC, abstractmethod
 
+# Configure logging for this module
+logger = logging.getLogger(__name__)
+# Ensure we show debug messages
+logging.basicConfig()
+logger.setLevel(logging.DEBUG)
+
 class Router(ABC):
     """Abstract base class for routing strategies"""
 
@@ -17,21 +24,44 @@ class RoundRobinRouter(Router):
     """Round-robin routing strategy"""
 
     def __init__(self, providers: List):
-        self.providers = providers
+        self.all_providers = providers
         self.index = 0
 
     def select(self, providers: List) -> Optional:
         if not providers:
+            logger.debug("Round-robin: No providers available")
             return None
 
-        # Find next available provider
-        for _ in range(len(providers)):
-            provider = self.providers[self.index % len(self.providers)]
-            self.index += 1
-            if provider in providers:
-                return provider
+        # If only one provider available, return it
+        if len(providers) == 1:
+            logger.debug(f"Round-robin: Only one provider: {providers[0].name}")
+            return providers[0]
 
-        return providers[0] if providers else None
+        logger.debug(f"Round-robin: Starting selection, index={self.index}, providers={[p.name for p in providers]}")
+        
+        # Find next available provider in round-robin fashion
+        start_index = self.index
+        attempts = 0
+        while attempts < len(self.all_providers):
+            # Get provider at current index from all providers
+            current_provider = self.all_providers[self.index % len(self.all_providers)]
+            next_index = (self.index + 1) % len(self.all_providers)
+            
+            logger.debug(f"Round-robin: Checking provider {current_provider.name} at index {self.index}")
+            
+            # Update index for next call
+            self.index = next_index
+            
+            # If this provider is in the available list, use it
+            if current_provider in providers:
+                logger.debug(f"Round-robin: Selected provider {current_provider.name}")
+                return current_provider
+            
+            attempts += 1
+        
+        # If we've cycled through all providers, just return first available
+        logger.debug(f"Round-robin: Fallback to first available: {providers[0].name}")
+        return providers[0]
 
 class WeightedRouter(Router):
     """Weighted random routing strategy"""

optillm/plugins/proxy_plugin.py

@@ -6,9 +6,9 @@
 """
 import logging
 from typing import Tuple, Optional
-from .proxy.config import ProxyConfig
-from .proxy.client import ProxyClient
-from .proxy.approach_handler import ApproachHandler
+from optillm.plugins.proxy.config import ProxyConfig
+from optillm.plugins.proxy.client import ProxyClient
+from optillm.plugins.proxy.approach_handler import ApproachHandler
 
 SLUG = "proxy"
 logger = logging.getLogger(__name__)
@@ -18,6 +18,9 @@
 log_level = os.environ.get('OPTILLM_LOG_LEVEL', 'INFO')
 logging.basicConfig(level=getattr(logging, log_level))
 
+# Global proxy client cache to maintain state between requests
+_proxy_client_cache = {}
+
 def run(system_prompt: str, initial_query: str, client, model: str, 
         request_config: dict = None) -> Tuple[str, int]:
     """
@@ -53,11 +56,18 @@ def run(system_prompt: str, initial_query: str, client, model: str,
             )
             return response.choices[0].message.content, response.usage.completion_tokens
 
-        # Create proxy client
-        proxy_client = ProxyClient(
-            config=config,
-            fallback_client=client
-        )
+        # Create or reuse proxy client to maintain state (important for round-robin)
+        config_key = str(config)  # Simple config-based cache key
+        if config_key not in _proxy_client_cache:
+            logger.debug("Creating new proxy client instance")
+            _proxy_client_cache[config_key] = ProxyClient(
+                config=config,
+                fallback_client=client
+            )
+        else:
+            logger.debug("Reusing existing proxy client instance")
+        
+        proxy_client = _proxy_client_cache[config_key]
 
         # Check for wrapped approach in extra_body (recommended method)
         wrapped_approach = None

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "mcp",
     "adaptive-classifier",
     "datasets",
+    "PyYAML",
     "selenium",
     "webdriver-manager",
     # MLX support for Apple Silicon optimization