Spaces:
Paused
refactor(usage): 🔨 cache provider plugin instances to reduce redundant instantiation
Browse filesIntroduced a provider instance cache (`_provider_instances`) to store and reuse provider plugin instances across multiple method calls.
- Added `_get_provider_instance()` helper method to centralize provider plugin instantiation logic with caching support
- Refactored `_get_usage_reset_config()`, `_get_model_quota_group()`, `_get_models_in_quota_group()`, `_get_usage_field_name()`, and cost calculation logic to use the cached provider instances
- Eliminated redundant provider plugin instantiation that occurred on every method call
- Simplified error handling by consolidating null checks in the helper method
This change improves performance by avoiding repeated instantiation of the same provider plugin objects and reduces code duplication across provider plugin access patterns.
|
@@ -76,6 +76,7 @@ class UsageManager:
|
|
| 76 |
self.rotation_tolerance = rotation_tolerance
|
| 77 |
self.provider_rotation_modes = provider_rotation_modes or {}
|
| 78 |
self.provider_plugins = provider_plugins or PROVIDER_PLUGINS
|
|
|
|
| 79 |
self.key_states: Dict[str, Dict[str, Any]] = {}
|
| 80 |
|
| 81 |
self._data_lock = asyncio.Lock()
|
|
@@ -138,6 +139,33 @@ class UsageManager:
|
|
| 138 |
|
| 139 |
return None
|
| 140 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 141 |
def _get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]:
|
| 142 |
"""
|
| 143 |
Get the usage reset configuration for a credential from its provider plugin.
|
|
@@ -150,15 +178,10 @@ class UsageManager:
|
|
| 150 |
or None to use default daily reset.
|
| 151 |
"""
|
| 152 |
provider = self._get_provider_from_credential(credential)
|
| 153 |
-
|
| 154 |
-
return None
|
| 155 |
|
| 156 |
-
|
| 157 |
-
|
| 158 |
-
return None
|
| 159 |
-
|
| 160 |
-
if hasattr(plugin, "get_usage_reset_config"):
|
| 161 |
-
return plugin.get_usage_reset_config(credential)
|
| 162 |
|
| 163 |
return None
|
| 164 |
|
|
@@ -187,15 +210,10 @@ class UsageManager:
|
|
| 187 |
Group name (e.g., "claude") or None if not grouped
|
| 188 |
"""
|
| 189 |
provider = self._get_provider_from_credential(credential)
|
| 190 |
-
|
| 191 |
-
return None
|
| 192 |
-
|
| 193 |
-
plugin = self.provider_plugins.get(provider)
|
| 194 |
-
if not plugin:
|
| 195 |
-
return None
|
| 196 |
|
| 197 |
-
if hasattr(
|
| 198 |
-
return
|
| 199 |
|
| 200 |
return None
|
| 201 |
|
|
@@ -211,15 +229,10 @@ class UsageManager:
|
|
| 211 |
List of full model names (e.g., ["antigravity/claude-opus-4-5", ...])
|
| 212 |
"""
|
| 213 |
provider = self._get_provider_from_credential(credential)
|
| 214 |
-
|
| 215 |
-
return []
|
| 216 |
-
|
| 217 |
-
plugin = self.provider_plugins.get(provider)
|
| 218 |
-
if not plugin:
|
| 219 |
-
return []
|
| 220 |
|
| 221 |
-
if hasattr(
|
| 222 |
-
models =
|
| 223 |
# Add provider prefix
|
| 224 |
return [f"{provider}/{m}" for m in models]
|
| 225 |
|
|
@@ -244,10 +257,10 @@ class UsageManager:
|
|
| 244 |
|
| 245 |
# Check provider default
|
| 246 |
provider = self._get_provider_from_credential(credential)
|
| 247 |
-
|
| 248 |
-
|
| 249 |
-
|
| 250 |
-
|
| 251 |
|
| 252 |
return "daily"
|
| 253 |
|
|
@@ -1302,10 +1315,10 @@ class UsageManager:
|
|
| 1302 |
)
|
| 1303 |
try:
|
| 1304 |
provider_name = model.split("/")[0]
|
| 1305 |
-
|
| 1306 |
|
| 1307 |
-
if
|
| 1308 |
-
|
| 1309 |
):
|
| 1310 |
lib_logger.debug(
|
| 1311 |
f"Skipping cost calculation for provider '{provider_name}' (custom provider)."
|
|
|
|
| 76 |
self.rotation_tolerance = rotation_tolerance
|
| 77 |
self.provider_rotation_modes = provider_rotation_modes or {}
|
| 78 |
self.provider_plugins = provider_plugins or PROVIDER_PLUGINS
|
| 79 |
+
self._provider_instances: Dict[str, Any] = {} # Cache for provider instances
|
| 80 |
self.key_states: Dict[str, Dict[str, Any]] = {}
|
| 81 |
|
| 82 |
self._data_lock = asyncio.Lock()
|
|
|
|
| 139 |
|
| 140 |
return None
|
| 141 |
|
| 142 |
+
def _get_provider_instance(self, provider: str) -> Optional[Any]:
|
| 143 |
+
"""
|
| 144 |
+
Get or create a provider plugin instance.
|
| 145 |
+
|
| 146 |
+
Args:
|
| 147 |
+
provider: The provider name
|
| 148 |
+
|
| 149 |
+
Returns:
|
| 150 |
+
Provider plugin instance or None
|
| 151 |
+
"""
|
| 152 |
+
if not provider:
|
| 153 |
+
return None
|
| 154 |
+
|
| 155 |
+
plugin_class = self.provider_plugins.get(provider)
|
| 156 |
+
if not plugin_class:
|
| 157 |
+
return None
|
| 158 |
+
|
| 159 |
+
# Get or create provider instance from cache
|
| 160 |
+
if provider not in self._provider_instances:
|
| 161 |
+
# Instantiate the plugin if it's a class, or use it directly if already an instance
|
| 162 |
+
if isinstance(plugin_class, type):
|
| 163 |
+
self._provider_instances[provider] = plugin_class()
|
| 164 |
+
else:
|
| 165 |
+
self._provider_instances[provider] = plugin_class
|
| 166 |
+
|
| 167 |
+
return self._provider_instances[provider]
|
| 168 |
+
|
| 169 |
def _get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]:
|
| 170 |
"""
|
| 171 |
Get the usage reset configuration for a credential from its provider plugin.
|
|
|
|
| 178 |
or None to use default daily reset.
|
| 179 |
"""
|
| 180 |
provider = self._get_provider_from_credential(credential)
|
| 181 |
+
plugin_instance = self._get_provider_instance(provider)
|
|
|
|
| 182 |
|
| 183 |
+
if plugin_instance and hasattr(plugin_instance, "get_usage_reset_config"):
|
| 184 |
+
return plugin_instance.get_usage_reset_config(credential)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 185 |
|
| 186 |
return None
|
| 187 |
|
|
|
|
| 210 |
Group name (e.g., "claude") or None if not grouped
|
| 211 |
"""
|
| 212 |
provider = self._get_provider_from_credential(credential)
|
| 213 |
+
plugin_instance = self._get_provider_instance(provider)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 214 |
|
| 215 |
+
if plugin_instance and hasattr(plugin_instance, "get_model_quota_group"):
|
| 216 |
+
return plugin_instance.get_model_quota_group(model)
|
| 217 |
|
| 218 |
return None
|
| 219 |
|
|
|
|
| 229 |
List of full model names (e.g., ["antigravity/claude-opus-4-5", ...])
|
| 230 |
"""
|
| 231 |
provider = self._get_provider_from_credential(credential)
|
| 232 |
+
plugin_instance = self._get_provider_instance(provider)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 233 |
|
| 234 |
+
if plugin_instance and hasattr(plugin_instance, "get_models_in_quota_group"):
|
| 235 |
+
models = plugin_instance.get_models_in_quota_group(group)
|
| 236 |
# Add provider prefix
|
| 237 |
return [f"{provider}/{m}" for m in models]
|
| 238 |
|
|
|
|
| 257 |
|
| 258 |
# Check provider default
|
| 259 |
provider = self._get_provider_from_credential(credential)
|
| 260 |
+
plugin_instance = self._get_provider_instance(provider)
|
| 261 |
+
|
| 262 |
+
if plugin_instance and hasattr(plugin_instance, "get_default_usage_field_name"):
|
| 263 |
+
return plugin_instance.get_default_usage_field_name()
|
| 264 |
|
| 265 |
return "daily"
|
| 266 |
|
|
|
|
| 1315 |
)
|
| 1316 |
try:
|
| 1317 |
provider_name = model.split("/")[0]
|
| 1318 |
+
provider_instance = self._get_provider_instance(provider_name)
|
| 1319 |
|
| 1320 |
+
if provider_instance and getattr(
|
| 1321 |
+
provider_instance, "skip_cost_calculation", False
|
| 1322 |
):
|
| 1323 |
lib_logger.debug(
|
| 1324 |
f"Skipping cost calculation for provider '{provider_name}' (custom provider)."
|