Mirrowel commited on
Commit
fd01482
·
1 Parent(s): 4bc7613

refactor(usage): 🔨 cache provider plugin instances to reduce redundant instantiation

Browse files

Introduced 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.

Files changed (1) hide show
  1. src/rotator_library/usage_manager.py +44 -31
src/rotator_library/usage_manager.py CHANGED
@@ -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
- if not provider:
154
- return None
155
 
156
- plugin = self.provider_plugins.get(provider)
157
- if not plugin:
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
- if not provider:
191
- return None
192
-
193
- plugin = self.provider_plugins.get(provider)
194
- if not plugin:
195
- return None
196
 
197
- if hasattr(plugin, "get_model_quota_group"):
198
- return plugin.get_model_quota_group(model)
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
- if not provider:
215
- return []
216
-
217
- plugin = self.provider_plugins.get(provider)
218
- if not plugin:
219
- return []
220
 
221
- if hasattr(plugin, "get_models_in_quota_group"):
222
- models = plugin.get_models_in_quota_group(group)
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
- if provider:
248
- plugin = self.provider_plugins.get(provider)
249
- if plugin and hasattr(plugin, "get_default_usage_field_name"):
250
- return plugin.get_default_usage_field_name()
251
 
252
  return "daily"
253
 
@@ -1302,10 +1315,10 @@ class UsageManager:
1302
  )
1303
  try:
1304
  provider_name = model.split("/")[0]
1305
- provider_plugin = self.provider_plugins.get(provider_name)
1306
 
1307
- if provider_plugin and getattr(
1308
- provider_plugin, "skip_cost_calculation", False
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)."