feat: Add Gemini API integration

agentops/llms/providers/gemini.py

@@ -0,0 +1,186 @@
+from typing import Optional, Generator, Any, Dict, Union
+
+from agentops.llms.providers.base import BaseProvider
+from agentops.event import LLMEvent, ErrorEvent
+from agentops.session import Session
+from agentops.helpers import get_ISO_time, check_call_stack_for_agent_id
+from agentops.log_config import logger
+from agentops.singleton import singleton
+
+# Store original methods at module level
+_ORIGINAL_METHODS = {}
+
+
+@singleton
+class GeminiProvider(BaseProvider):
+    """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.
+
+        Args:
+            client: Optional client instance. If not provided, will be set during override.
+        """
+        super().__init__(client)
+        self._provider_name = "Gemini"
+
+    def handle_response(
+        self, response, kwargs, init_timestamp, session: Optional[Session] = None
+    ) -> Union[Any, Generator[Any, None, None]]:
+        """Handle responses from Gemini API for both sync and streaming modes.
+
+        Args:
+            response: The response from the Gemini API
+            kwargs: The keyword arguments passed to generate_content
+            init_timestamp: The timestamp when the request was initiated
+            session: Optional AgentOps session for recording events
+
+        Returns:
+            For sync responses: The original response object
+            For streaming responses: A generator yielding response chunks
+
+        Note:
+            Token counts are extracted from usage_metadata if available.
+        """
+        llm_event = LLMEvent(init_timestamp=init_timestamp, params=kwargs)
+        if session is not None:
+            llm_event.session_id = session.session_id
+
+        # For streaming responses
+        if kwargs.get("stream", False):
+            accumulated_text = []  # Use list to accumulate text chunks
+
+            def handle_stream_chunk(chunk):
+                nonlocal llm_event
+                try:
+                    if llm_event.returns is None:
+                        llm_event.returns = chunk
+                        llm_event.agent_id = check_call_stack_for_agent_id()
+                        llm_event.model = getattr(chunk, "model", "gemini-1.5-flash")
+                        llm_event.prompt = kwargs.get("prompt", kwargs.get("contents", []))
+
+                    if hasattr(chunk, "text") and chunk.text:
+                        accumulated_text.append(chunk.text)
+
+                    # Extract token counts if available
+                    if hasattr(chunk, "usage_metadata"):
+                        usage = chunk.usage_metadata
+                        llm_event.prompt_tokens = getattr(usage, "prompt_token_count", None)
+                        llm_event.completion_tokens = getattr(usage, "candidates_token_count", None)
+
+                    # If this is the last chunk
+                    if hasattr(chunk, "finish_reason") and chunk.finish_reason:
+                        llm_event.completion = "".join(accumulated_text)
+                        llm_event.end_timestamp = get_ISO_time()
+                        self._safe_record(session, llm_event)
+
+                except Exception as e:
+                    if session is not None:
+                        self._safe_record(session, ErrorEvent(trigger_event=llm_event, exception=e))
+                    logger.warning(
+                        f"Unable to parse chunk for Gemini LLM call. Error: {str(e)}\n"
+                        f"Chunk: {chunk}\n"
+                        f"kwargs: {kwargs}\n"
+                    )
+
+            def stream_handler(stream):
+                try:
+                    for chunk in stream:
+                        handle_stream_chunk(chunk)
+                        yield chunk
+                except Exception as e:
+                    if session is not None:
+                        self._safe_record(session, ErrorEvent(trigger_event=llm_event, exception=e))
+                    raise  # Re-raise after recording error
+
+            return stream_handler(response)
+
+        # For synchronous responses
+        try:
+            llm_event.returns = response
+            llm_event.agent_id = check_call_stack_for_agent_id()
+            llm_event.prompt = kwargs.get("prompt", kwargs.get("contents", []))
+            llm_event.completion = response.text
+            llm_event.model = getattr(response, "model", "gemini-1.5-flash")
+
+            # Extract token counts from usage metadata if available
+            if hasattr(response, "usage_metadata"):
+                usage = response.usage_metadata
+                llm_event.prompt_tokens = getattr(usage, "prompt_token_count", None)
+                llm_event.completion_tokens = getattr(usage, "candidates_token_count", None)
+
+            llm_event.end_timestamp = get_ISO_time()
+            self._safe_record(session, llm_event)
+        except Exception as e:
+            if session is not None:
+                self._safe_record(session, ErrorEvent(trigger_event=llm_event, exception=e))
+            logger.warning(
+                f"Unable to parse response for Gemini LLM call. Error: {str(e)}\n"
+                f"Response: {response}\n"
+                f"kwargs: {kwargs}\n"
+            )
+
+        return response
+
+    def override(self):
+        """Override Gemini's generate_content method to track LLM events.
+
+        Note:
+            This method is called automatically by AgentOps during initialization.
+            Users should not call this method directly."""
+        import google.generativeai as genai
+
+        # Store original method if not already stored
+        if "generate_content" not in _ORIGINAL_METHODS:
+            _ORIGINAL_METHODS["generate_content"] = genai.GenerativeModel.generate_content
+
+        # Store provider instance for the closure
+        provider = self
+
+        def patched_function(self, *args, **kwargs):
+            init_timestamp = get_ISO_time()
+
+            # Extract and remove session from kwargs if present
+            session = kwargs.pop("session", None)
+
+            # 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
+                prompt = args[0]
+                if "contents" not in kwargs:
+                    kwargs["contents"] = prompt
+                    event_kwargs["prompt"] = prompt  # Store original prompt for event tracking
+                args = args[1:]  # Remove prompt from args since we moved it to kwargs
+
+            # Call original method and track event
+            try:
+                if "generate_content" in _ORIGINAL_METHODS:
+                    result = _ORIGINAL_METHODS["generate_content"](self, *args, **kwargs)
+                    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
+            except Exception as e:
+                logger.error(f"Error in Gemini generate_content: {str(e)}")
+                if session is not None:
+                    provider._safe_record(session, ErrorEvent(exception=e))
+                raise  # Re-raise the exception after recording
+
+        # Override the method at class level
+        genai.GenerativeModel.generate_content = patched_function
+
+    def undo_override(self):
+        """Restore original Gemini methods.
+
+        Note:
+            This method is called automatically by AgentOps during cleanup.
+            Users should not call this method directly."""
+        if "generate_content" in _ORIGINAL_METHODS:
+            import google.generativeai as genai
+
+            genai.GenerativeModel.generate_content = _ORIGINAL_METHODS["generate_content"]

agentops/llms/tracker.py

@@ -16,6 +16,7 @@
 from .providers.ai21 import AI21Provider
 from .providers.llama_stack_client import LlamaStackClientProvider
 from .providers.taskweaver import TaskWeaverProvider
+from .providers.gemini import GeminiProvider
 
 original_func = {}
 original_create = None
@@ -24,6 +25,9 @@
 
 class LlmTracker:
     SUPPORTED_APIS = {
+        "google.generativeai": {
+            "0.1.0": ("GenerativeModel.generate_content", "GenerativeModel.generate_content_stream"),
+        },
         "litellm": {"1.3.1": ("openai_chat_completions.completion",)},
         "openai": {
             "1.0.0": (
@@ -210,6 +214,26 @@
                     else:
                         logger.warning(f"Only TaskWeaver>=0.0.1 supported. v{module_version} found.")
 
+                if api == "google.generativeai":
+                    module_version = version(api)
+
+                    if Version(module_version) >= parse("0.1.0"):
+                        import google.generativeai as genai
+                        import os
+
+                        api_key = os.getenv("GEMINI_API_KEY")
+                        if api_key:
+                            try:
+                                genai.configure(api_key=api_key)
+                                provider = GeminiProvider(self.client)
+                                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:
+                        logger.warning(f"Only google.generativeai>=0.1.0 supported. v{module_version} found.")
+
     def stop_instrumenting(self):
         OpenAiProvider(self.client).undo_override()
         GroqProvider(self.client).undo_override()
@@ -221,3 +245,4 @@
         AI21Provider(self.client).undo_override()
         LlamaStackClientProvider(self.client).undo_override()
         TaskWeaverProvider(self.client).undo_override()
+        GeminiProvider(self.client).undo_override()

examples/gemini_examples/gemini_example_sync.ipynb

@@ -0,0 +1,135 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "580c85ac",
+   "metadata": {},
+   "source": [
+    "# Gemini API Example with AgentOps\n",
+    "\n",
+    "This notebook demonstrates how to use AgentOps with Google's Gemini API for both synchronous and streaming text generation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d731924a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import google.generativeai as genai\n",
+    "import agentops\n",
+    "from agentops.llms.providers.gemini import GeminiProvider"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a94545c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Configure the Gemini API\n",
+    "import os\n",
+    "\n",
+    "# Replace with your API key\n",
+    "# You can get one at: https://ai.google.dev/tutorials/setup\n",
+    "GEMINI_API_KEY = \"YOUR_API_KEY_HERE\"  # Replace with your API key\n",
+    "genai.configure(api_key=GEMINI_API_KEY)\n",
+    "\n",
+    "# Note: In production, use environment variables:\n",
+    "# import os\n",
+    "# GEMINI_API_KEY = os.getenv(\"GEMINI_API_KEY\")\n",
+    "# genai.configure(api_key=GEMINI_API_KEY)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d632fe48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize AgentOps and Gemini model\n",
+    "ao_client = agentops.init()\n",
+    "model = genai.GenerativeModel(\"gemini-1.5-flash\")\n",
+    "\n",
+    "# Initialize and override Gemini provider\n",
+    "provider = GeminiProvider(model)\n",
+    "provider.override()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3923b6b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test synchronous generation\n",
+    "print(\"Testing synchronous generation:\")\n",
+    "response = model.generate_content(\n",
+    "    \"What are the three laws of robotics?\",\n",
+    "    session=ao_client\n",
+    ")\n",
+    "print(response.text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da54e521",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test streaming generation\n",
+    "print(\"\\nTesting streaming generation:\")\n",
+    "response = model.generate_content(\n",
+    "    \"Explain the concept of machine learning in simple terms.\",\n",
+    "    stream=True,\n",
+    "    session=ao_client\n",
+    ")\n",
+    "\n",
+    "for chunk in response:\n",
+    "    print(chunk.text, end=\"\")\n",
+    "print()  # Add newline after streaming output\n",
+    "\n",
+    "# Test another synchronous generation\n",
+    "print(\"\\nTesting another synchronous generation:\")\n",
+    "response = model.generate_content(\n",
+    "    \"What is the difference between supervised and unsupervised learning?\",\n",
+    "    session=ao_client\n",
+    ")\n",
+    "print(response.text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6a674c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# End session and check stats\n",
+    "agentops.end_session(\n",
+    "    end_state=\"Success\",\n",
+    "    end_state_reason=\"Gemini integration example completed successfully\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b6d35f28",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Clean up\n",
+    "provider.undo_override()"
+   ]
+  }
+ ],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}