feat: Add Gemini API integration

agentops/llms/providers/gemini.py

@@ -0,0 +1,194 @@
+from typing import Optional, 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
+
+
+@singleton
+class GeminiProvider(BaseProvider):
+    original_generate_content = None
+    original_generate_content_async = None
+
+    """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) -> dict:
+        """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
+        """
+        llm_event = LLMEvent(init_timestamp=init_timestamp, params=kwargs)
+        if session is not None:
+            llm_event.session_id = session.session_id
+
+        accumulated_content = ""
+
+        def handle_stream_chunk(chunk):
+            nonlocal llm_event, accumulated_content
+            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", None) or "gemini-1.5-flash"
+                    llm_event.prompt = kwargs.get("prompt", kwargs.get("contents", None)) or []
+
+                # Accumulate text from chunk
+                if hasattr(chunk, "text") and chunk.text:
+                    accumulated_content += chunk.text
+
+                # Extract token counts if available
+                if hasattr(chunk, "usage_metadata"):
+                    llm_event.prompt_tokens = getattr(chunk.usage_metadata, "prompt_token_count", None)
+                    llm_event.completion_tokens = getattr(chunk.usage_metadata, "candidates_token_count", None)
+
+                # If this is the last chunk
+                if hasattr(chunk, "finish_reason") and chunk.finish_reason:
+                    llm_event.completion = accumulated_content
+                    llm_event.end_timestamp = get_ISO_time()
+                    self._safe_record(session, llm_event)
+
+            except Exception as e:
+                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"Response: {chunk}\n"
+                    f"Arguments: {kwargs}\n"
+                )
+
+        # For streaming responses
+        if kwargs.get("stream", False):
+
+            def generator():
+                for chunk in response:
+                    handle_stream_chunk(chunk)
+                    yield chunk
+
+            return generator()
+
+        # 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", None)) or []
+            llm_event.completion = response.text
+            llm_event.model = getattr(response, "model", None) or "gemini-1.5-flash"
+
+            # Extract token counts from usage metadata if available
+            if hasattr(response, "usage_metadata"):
+                llm_event.prompt_tokens = getattr(response.usage_metadata, "prompt_token_count", None)
+                llm_event.completion_tokens = getattr(response.usage_metadata, "candidates_token_count", None)
+
+            llm_event.end_timestamp = get_ISO_time()
+            self._safe_record(session, llm_event)
+        except Exception as e:
+            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"Arguments: {kwargs}\n"
+            )
+
+        return response
+
+    def override(self):
+        """Override Gemini's generate_content method to track LLM events."""
+        self._override_gemini_generate_content()
+        self._override_gemini_generate_content_async()
+
+    def _override_gemini_generate_content(self):
+        """Override synchronous generate_content method"""
+        import google.generativeai as genai
+
+        # Store original method if not already stored
+        if self.original_generate_content is None:
+            self.original_generate_content = genai.GenerativeModel.generate_content
+
+        provider = self  # Store provider instance for closure
+
+        def patched_function(model_self, *args, **kwargs):
+            init_timestamp = get_ISO_time()
+            session = kwargs.pop("session", None)
+
+            # Handle positional prompt argument
+            event_kwargs = kwargs.copy()
+            if args and len(args) > 0:
+                prompt = args[0]
+                if "contents" not in kwargs:
+                    kwargs["contents"] = prompt
+                    event_kwargs["prompt"] = prompt
+                args = args[1:]
+
+            result = provider.original_generate_content(model_self, *args, **kwargs)
+            return provider.handle_response(result, event_kwargs, init_timestamp, session=session)
+
+        # Override the method at class level
+        genai.GenerativeModel.generate_content = patched_function
+
+    def _override_gemini_generate_content_async(self):
+        """Override asynchronous generate_content method"""
+        import google.generativeai as genai
+
+        # Store original async method if not already stored
+        if self.original_generate_content_async is None:
+            self.original_generate_content_async = genai.GenerativeModel.generate_content_async
+
+        provider = self  # Store provider instance for closure
+
+        async def patched_function(model_self, *args, **kwargs):
+            init_timestamp = get_ISO_time()
+            session = kwargs.pop("session", None)
+
+            # Handle positional prompt argument
+            event_kwargs = kwargs.copy()
+            if args and len(args) > 0:
+                prompt = args[0]
+                if "contents" not in kwargs:
+                    kwargs["contents"] = prompt
+                    event_kwargs["prompt"] = prompt
+                args = args[1:]
+
+            result = await provider.original_generate_content_async(model_self, *args, **kwargs)
+            return provider.handle_response(result, event_kwargs, init_timestamp, session=session)
+
+        # Override the async method at class level
+        genai.GenerativeModel.generate_content_async = 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."""
+        import google.generativeai as genai
+
+        if self.original_generate_content is not None:
+            genai.GenerativeModel.generate_content = self.original_generate_content
+            self.original_generate_content = None
+
+        if self.original_generate_content_async is not None:
+            genai.GenerativeModel.generate_content_async = self.original_generate_content_async
+            self.original_generate_content_async = None

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,15 @@
                     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"):
+                        provider = GeminiProvider(self.client)
+                        provider.override()
+                    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 +234,4 @@
         AI21Provider(self.client).undo_override()
         LlamaStackClientProvider(self.client).undo_override()
         TaskWeaverProvider(self.client).undo_override()
+        GeminiProvider(self.client).undo_override()

docs/mint.json

@@ -93,6 +93,7 @@
         "v1/integrations/camel",
         "v1/integrations/cohere",
         "v1/integrations/crewai",
+        "v1/integrations/gemini",
         "v1/integrations/groq",
         "v1/integrations/langchain",
         "v1/integrations/llama_stack",

docs/v1/examples/examples.mdx

@@ -57,6 +57,10 @@ mode: "wide"
     Ultra-fast LLM inference with Groq Cloud
   </Card>
 
+  <Card title="Gemini" icon={<img src="https://www.github.com/agentops-ai/agentops/blob/main/docs/images/external/deepmind/gemini-logo.png?raw=true" alt="Gemini" />} iconType="image" href="/v1/integrations/gemini">
+    Explore Google DeepMind's Gemini with observation via AgentOps
+  </Card>
+
   <Card title="LangChain" icon={<img src="https://www.github.com/agentops-ai/agentops/blob/main/docs/images/external/langchain/langchain-logo.png?raw=true" alt="LangChain" />} iconType="image" href="/v1/examples/langchain">
     Jupyter Notebook with a sample LangChain integration
   </Card>

docs/v1/integrations/gemini.mdx

@@ -0,0 +1,118 @@
+---
+title: Gemini
+description: "Explore Google DeepMind's Gemini with observation via AgentOps"
+---
+
+import CodeTooltip from '/snippets/add-code-tooltip.mdx'
+import EnvTooltip from '/snippets/add-env-tooltip.mdx'
+
+[Gemini (Google Generative AI)](https://ai.google.dev/gemini-api/docs/quickstart) is a leading provider of AI tools and services.
+Explore the [Gemini API](https://ai.google.dev/docs) for more information.
+
+<Note>
+  `google-generativeai>=0.1.0` is currently supported.
+</Note>
+
+<Steps>
+  <Step title="Install the AgentOps SDK">
+    <CodeGroup>
+      ```bash pip 
+      pip install agentops
+      ```
+      ```bash poetry
+      poetry add agentops
+      ```
+    </CodeGroup>
+  </Step>
+  <Step title="Install the Gemini SDK">
+    <Note>
+      `google-generativeai>=0.1.0` is required for Gemini integration.
+    </Note>
+    <CodeGroup>
+      ```bash pip
+      pip install google-generativeai
+      ```
+      ```bash poetry
+      poetry add google-generativeai
+      ```
+    </CodeGroup>
+  </Step>
+  <Step title="Add 3 lines of code">
+    <CodeTooltip/>
+    <CodeGroup>
+      ```python python
+      import google.generativeai as genai
+      import agentops
+
+      agentops.init(<INSERT YOUR API KEY HERE>)
+      model = genai.GenerativeModel("gemini-1.5-flash")
+      ...
+      # End of program (e.g. main.py)
+      agentops.end_session("Success") # Success|Fail|Indeterminate
+      ```
+    </CodeGroup>
+    <EnvTooltip />
+      <CodeGroup>
+        ```python .env
+        AGENTOPS_API_KEY=<YOUR API KEY>
+        GEMINI_API_KEY=<YOUR GEMINI API KEY>
+        ```
+      </CodeGroup>
+      Read more about environment variables in [Advanced Configuration](/v1/usage/advanced-configuration)
+  </Step>
+  <Step title="Run your Agent">
+    Execute your program and visit [app.agentops.ai/drilldown](https://app.agentops.ai/drilldown) to observe your Agent! 🕵️
+    <Tip>
+      After your run, AgentOps prints a clickable url to console linking directly to your session in the Dashboard
+    </Tip>
+    <div/>
+    <Frame type="glass" caption="Clickable link to session">
+      <img height="200" src="https://github.com/AgentOps-AI/agentops/blob/main/docs/images/link-to-session.gif?raw=true" />
+    </Frame>
+  </Step>
+</Steps>
+
+## Full Examples
+
+<CodeGroup>
+  ```python sync
+  import google.generativeai as genai
+  import agentops
+
+  agentops.init(<INSERT YOUR API KEY HERE>)
+  model = genai.GenerativeModel("gemini-1.5-flash")
+
+  response = model.generate_content(
+      "Write a haiku about AI and humans working together"
+  )
+
+  print(response.text)
+  agentops.end_session('Success')
+  ```
+
+  ```python stream
+  import google.generativeai as genai
+  import agentops
+
+  agentops.init(<INSERT YOUR API KEY HERE>)
+  model = genai.GenerativeModel("gemini-1.5-flash")
+
+  response = model.generate_content(
+      "Write a haiku about AI and humans working together",
+      stream=True
+  )
+
+  for chunk in response:
+      print(chunk.text, end="")
+
+  agentops.end_session('Success')
+  ```
+</CodeGroup>
+
+You can find more examples in the [Gemini Examples](/v1/examples/gemini_examples) section.
+
+<script type="module" src="/scripts/github_stars.js"></script>
+<script type="module" src="/scripts/scroll-img-fadein-animation.js"></script>
+<script type="module" src="/scripts/button_heartbeat_animation.js"></script>
+<script type="css" src="/styles/styles.css"></script>
+<script type="module" src="/scripts/adjust_api_dynamically.js"></script>