fxies

Author: codelion

optillm/plugins/proxy/client.py

@@ -47,6 +47,13 @@ def client(self):
                     api_version="2024-02-01",
                     max_retries=0  # Disable client retries - we handle them
                 )
+            elif 'generativelanguage.googleapis.com' in self.base_url:
+                # Google AI client - create custom client to avoid "models/" prefix
+                from optillm.plugins.proxy.google_client import GoogleAIClient
+                self._client = GoogleAIClient(
+                    api_key=self.api_key,
+                    base_url=self.base_url
+                )
             else:
                 # Standard OpenAI-compatible client
                 self._client = OpenAI(
@@ -215,8 +222,8 @@ def create(self, **kwargs):
 
                     attempted_providers.add(provider)
 
-                    # Try to acquire a slot for this provider (with short timeout to try next provider quickly)
-                    slot_timeout = 0.5  # Don't wait too long for a single provider
+                    # Try to acquire a slot for this provider (with reasonable timeout for queueing)
+                    slot_timeout = 10.0  # Wait up to 10 seconds for provider to become available
                     if not provider.acquire_slot(timeout=slot_timeout):
                         logger.debug(f"Provider {provider.name} at max capacity, trying next provider")
                         errors.append((provider.name, "At max concurrent requests"))

optillm/plugins/proxy/google_client.py

@@ -0,0 +1,92 @@
+"""
+Custom Google AI client that doesn't add "models/" prefix to model names
+"""
+import requests
+import json
+from typing import Dict, List, Any
+
+
+class GoogleAIClient:
+    """Custom client for Google AI that bypasses OpenAI client's model name prefix behavior"""
+    
+    def __init__(self, api_key: str, base_url: str):
+        self.api_key = api_key
+        self.base_url = base_url.rstrip('/')
+        self.chat = self.Chat(self)
+        self.models = self.Models(self)
+    
+    class Chat:
+        def __init__(self, client):
+            self.client = client
+            self.completions = self.Completions(client)
+        
+        class Completions:
+            def __init__(self, client):
+                self.client = client
+            
+            def create(self, model: str, messages: List[Dict[str, str]], **kwargs) -> Any:
+                """Create chat completion without adding models/ prefix to model name"""
+                url = f"{self.client.base_url}/chat/completions"
+                
+                headers = {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {self.client.api_key}"
+                }
+                
+                # Build request data - use model name directly without "models/" prefix
+                data = {
+                    "model": model,  # Use exactly as provided - no prefix!
+                    "messages": messages,
+                    **kwargs
+                }
+                
+                # Make direct HTTP request to bypass OpenAI client behavior
+                response = requests.post(url, headers=headers, json=data, timeout=kwargs.get('timeout', 30))
+                
+                if response.status_code != 200:
+                    error_text = response.text
+                    raise Exception(f"HTTP {response.status_code}: {error_text}")
+                
+                # Parse response and return OpenAI-compatible object
+                result = response.json()
+                
+                # Create a simple object that has the attributes expected by the proxy
+                class CompletionResponse:
+                    def __init__(self, data):
+                        self._data = data
+                        self.choices = data.get('choices', [])
+                        self.usage = data.get('usage', {})
+                        self.model = data.get('model', model)
+                    
+                    def model_dump(self):
+                        return self._data
+                    
+                    def __getitem__(self, key):
+                        return self._data[key]
+                    
+                    def get(self, key, default=None):
+                        return self._data.get(key, default)
+                
+                return CompletionResponse(result)
+    
+    class Models:
+        def __init__(self, client):
+            self.client = client
+        
+        def list(self):
+            """Simple models list for health checking"""
+            url = f"{self.client.base_url}/models"
+            headers = {
+                "Authorization": f"Bearer {self.client.api_key}"
+            }
+            
+            try:
+                response = requests.get(url, headers=headers, timeout=5)
+                if response.status_code == 200:
+                    return response.json()
+                else:
+                    # Return a mock response if health check fails
+                    return {"data": [{"id": "gemma-3-4b-it"}]}
+            except:
+                # Return a mock response if health check fails
+                return {"data": [{"id": "gemma-3-4b-it"}]}