fix: Update Gemini provider to properly handle prompt extraction and improve test coverage

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

Author: devin-ai-integration[bot]

agentops/llms/providers/gemini.py

@@ -19,12 +19,6 @@ class GeminiProvider(BaseProvider):
     is called and the google.generativeai package is imported. No manual
     initialization is required."""
 
-    """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.
 
@@ -66,7 +60,7 @@ def handle_stream_chunk(chunk):
                     llm_event.returns = chunk
                     llm_event.agent_id = check_call_stack_for_agent_id()
                     llm_event.model = getattr(chunk, "model", "gemini-1.5-flash")  # Default if not provided
-                    llm_event.prompt = kwargs.get("contents", [])
+                    llm_event.prompt = kwargs.get("prompt") or kwargs.get("contents", [])
 
                 try:
                     if hasattr(chunk, "text") and chunk.text:
@@ -103,7 +97,7 @@ def stream_handler(stream):
         try:
             llm_event.returns = response
             llm_event.agent_id = check_call_stack_for_agent_id()
-            llm_event.prompt = kwargs.get("contents", [])
+            llm_event.prompt = kwargs.get("prompt") or kwargs.get("contents", [])
             llm_event.completion = response.text
             llm_event.model = getattr(response, "model", "gemini-1.5-flash")
 
@@ -144,10 +138,19 @@ def patched_function(self, *args, **kwargs):
             init_timestamp = get_ISO_time()
             session = kwargs.pop("session", None)  # Always try to pop session, returns None if not present
 
+            # Handle positional prompt argument
+            event_kwargs = kwargs.copy()  # Create a copy for event tracking
+            if args and len(args) > 0:
+                # First argument is the prompt
+                if "contents" not in kwargs:
+                    kwargs["contents"] = args[0]
+                    event_kwargs["prompt"] = args[0]  # Store original prompt
+                args = args[1:]  # Remove prompt from args since we moved it to kwargs
+
             # Call original method and track event
             if "generate_content" in _ORIGINAL_METHODS:
                 result = _ORIGINAL_METHODS["generate_content"](self, *args, **kwargs)
-                return provider.handle_response(result, kwargs, init_timestamp, session=session)
+                return provider.handle_response(result, event_kwargs, init_timestamp, session=session)
             else:
                 logger.error("Original generate_content method not found. Cannot proceed with override.")
                 return None

tests/unit/test_llms/providers/test_gemini.py

@@ -1,6 +1,6 @@
 import google.generativeai as genai
 import agentops
-from agentops.llms.providers.gemini import GeminiProvider
+from agentops.llms.providers.gemini import GeminiProvider, _ORIGINAL_METHODS
 from agentops.event import LLMEvent
 
 # Configure the API key from environment variable
@@ -43,7 +43,7 @@ def test_gemini_provider():
     provider = GeminiProvider(model)
     assert provider.client == model
     assert provider.provider_name == "Gemini"
-    assert provider.original_generate is None
+    assert "generate_content" not in _ORIGINAL_METHODS
 
 
 def test_gemini_sync_generation():
@@ -97,31 +97,31 @@ class InvalidClient:
 
     # Test override with None client
     provider.override()  # Should log warning and return
-    assert provider.original_generate is None
+    assert "generate_content" not in _ORIGINAL_METHODS
 
     # Test override with uninitialized generate_content
     provider.client = InvalidClient()
     provider.override()  # Should log warning about missing generate_content
-    assert provider.original_generate is None
+    assert "generate_content" not in _ORIGINAL_METHODS
 
-    # Test patched function with None original_generate
+    # Test patched function with missing original method
     model = genai.GenerativeModel("gemini-1.5-flash")
     provider = GeminiProvider(model)
-    provider.original_generate = None
     provider.override()
 
-    # Should log error and return None
+    # Should log error and return None when original method is missing
+    if "generate_content" in _ORIGINAL_METHODS:
+        del _ORIGINAL_METHODS["generate_content"]
     result = model.generate_content("test prompt")
     assert result is None
 
     # Test undo_override with None client
     provider.client = None
     provider.undo_override()  # Should handle None client gracefully
 
-    # Test undo_override with None original_generate
+    # Test undo_override with missing original method
     provider.client = model
-    provider.original_generate = None
-    provider.undo_override()  # Should handle None original_generate gracefully
+    provider.undo_override()  # Should handle missing original method gracefully
 
     # Test automatic provider detection
     agentops.init()
@@ -503,6 +503,7 @@ def test_undo_override():
     provider.undo_override()
     assert model.generate_content == original_generate
 
-    # Test undo_override when original_generate is None
-    provider.original_generate = None
+    # Test undo_override with missing original method
+    if "generate_content" in _ORIGINAL_METHODS:
+        del _ORIGINAL_METHODS["generate_content"]
     provider.undo_override()  # Should not raise any errors