fix: Use module-level storage for original method

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

Author: devin-ai-integration[bot]

agentops/llms/providers/gemini.py

@@ -7,6 +7,9 @@
 from agentops.log_config import logger
 from agentops.singleton import singleton
 
+# Store original methods at module level
+_ORIGINAL_METHODS = {}
+
 
 @singleton
 class GeminiProvider(BaseProvider):
@@ -16,7 +19,11 @@ class GeminiProvider(BaseProvider):
     is called and the google.generativeai package is imported. No manual
     initialization is required."""
 
-    _original_generate = None  # Store as class attribute
+    """Provider for Google's Gemini API.
+
+    This provider is automatically detected and initialized when agentops.init()
+    is called and the google.generativeai package is imported. No manual
+    initialization is required."""
 
     def __init__(self, client=None):
         """Initialize the Gemini provider.
@@ -127,8 +134,8 @@ def override(self):
         import google.generativeai as genai
 
         # Store original method if not already stored
-        if GeminiProvider._original_generate is None:
-            GeminiProvider._original_generate = genai.GenerativeModel.generate_content
+        if 'generate_content' not in _ORIGINAL_METHODS:
+            _ORIGINAL_METHODS['generate_content'] = genai.GenerativeModel.generate_content
 
         # Store provider instance for the closure
         provider = self
@@ -138,8 +145,8 @@ def patched_function(self, *args, **kwargs):
             session = kwargs.pop("session", None)  # Always try to pop session, returns None if not present
 
             # Call original method and track event
-            if GeminiProvider._original_generate:
-                result = GeminiProvider._original_generate(self, *args, **kwargs)
+            if 'generate_content' in _ORIGINAL_METHODS:
+                result = _ORIGINAL_METHODS['generate_content'](self, *args, **kwargs)
                 return provider.handle_response(result, kwargs, init_timestamp, session=session)
             else:
                 logger.error("Original generate_content method not found. Cannot proceed with override.")
@@ -154,6 +161,6 @@ def undo_override(self):
         Note:
             This method is called automatically by AgentOps during cleanup.
             Users should not call this method directly."""
-        if GeminiProvider._original_generate is not None:
+        if 'generate_content' in _ORIGINAL_METHODS:
             import google.generativeai as genai
-            genai.GenerativeModel.generate_content = GeminiProvider._original_generate
+            genai.GenerativeModel.generate_content = _ORIGINAL_METHODS['generate_content']

agentops/llms/tracker.py

@@ -222,10 +222,12 @@ def override_api(self):
                         import os
                         api_key = os.getenv("GEMINI_API_KEY")
                         if api_key:
-                            genai.configure(api_key=api_key)
-                            model = genai.GenerativeModel("gemini-1.5-flash")
-                            provider = GeminiProvider()
-                            provider.override()
+                            try:
+                                genai.configure(api_key=api_key)
+                                provider = GeminiProvider()
+                                provider.override()
+                            except Exception as e:
+                                logger.warning(f"Failed to initialize Gemini provider: {str(e)}")
                         else:
                             logger.warning("GEMINI_API_KEY environment variable is required for Gemini integration")
                     else: