From 420f826e561d4d5aa2fb98277f70fb0e6886ffab Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 12 Feb 2025 06:24:05 +0000
Subject: [PATCH] fix: Initialize agentops in Crew setup (#2102)

- Add proper agentops initialization in Crew's set_private_attrs
- Add test coverage for agentops initialization
- Add test coverage for Gemini LLM with agentops

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/crew.py | 15 ++++++++
 tests/crew_test.py | 85 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/src/crewai/crew.py b/src/crewai/crew.py
index 9ae9ce2c0..f72fac30e 100644
--- a/src/crewai/crew.py
+++ b/src/crewai/crew.py
@@ -1,5 +1,6 @@
 import asyncio
 import json
+import os
 import re
 import uuid
 import warnings
@@ -248,6 +249,20 @@ class Crew(BaseModel):
         if self.output_log_file:
             self._file_handler = FileHandler(self.output_log_file)
         self._rpm_controller = RPMController(max_rpm=self.max_rpm, logger=self._logger)
+
+        # Initialize agentops if available and API key is present
+        if agentops:
+            api_key = os.getenv("AGENTOPS_API_KEY")
+            if api_key:
+                try:
+                    agentops.init(api_key)
+                except Exception as e:
+                    self._logger.log(
+                        "warning",
+                        f"Failed to initialize agentops: {e}",
+                        color="yellow"
+                    )
+
         if self.function_calling_llm and not isinstance(self.function_calling_llm, LLM):
             self.function_calling_llm = create_llm(self.function_calling_llm)
 
diff --git a/tests/crew_test.py b/tests/crew_test.py
index 0539ea347..b773d9c69 100644
--- a/tests/crew_test.py
+++ b/tests/crew_test.py
@@ -27,6 +27,91 @@ from crewai.utilities import Logger
 from crewai.utilities.rpm_controller import RPMController
 from crewai.utilities.task_output_storage_handler import TaskOutputStorageHandler
 
+def test_agentops_initialization_with_api_key(monkeypatch):
+    """Test that agentops is properly initialized when API key is present."""
+    import agentops
+    
+    # Mock agentops.init
+    mock_init = MagicMock()
+    monkeypatch.setattr(agentops, "init", mock_init)
+    
+    # Set API key
+    monkeypatch.setenv("AGENTOPS_API_KEY", "test-key")
+    
+    # Create crew
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher,
+    )
+    crew = Crew(agents=[researcher], tasks=[task])
+    
+    # Verify agentops was initialized
+    mock_init.assert_called_once_with("test-key")
+
+def test_agentops_initialization_without_api_key(monkeypatch):
+    """Test that agentops is not initialized when API key is not present."""
+    import agentops
+    
+    # Mock agentops.init
+    mock_init = MagicMock()
+    monkeypatch.setattr(agentops, "init", mock_init)
+    
+    # Create crew without setting API key
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher,
+    )
+    crew = Crew(agents=[researcher], tasks=[task])
+    
+    # Verify agentops.init was not called
+    mock_init.assert_not_called()
+
+def test_gemini_llm_with_agentops(monkeypatch):
+    """Test that Gemini LLM works correctly with agentops."""
+    from crewai.llm import LLM
+    import agentops
+    
+    # Mock agentops
+    mock_agentops = MagicMock()
+    monkeypatch.setattr("crewai.crew.agentops", mock_agentops)
+    
+    # Set API keys
+    monkeypatch.setenv("AGENTOPS_API_KEY", "test-key")
+    monkeypatch.setenv("GOOGLE_API_KEY", "test-key")
+    
+    # Create crew with Gemini LLM
+    llm = LLM(model="gemini-pro")
+    agent = Agent(
+        role="test",
+        goal="test",
+        backstory="test",
+        llm=llm
+    )
+    task = Task(
+        description="test task",
+        expected_output="test output",
+        agent=agent
+    )
+    crew = Crew(agents=[agent], tasks=[task])
+    
+    # Mock the agent execution to avoid actual API calls
+    with patch.object(Task, 'execute_sync', return_value=TaskOutput(
+        description="test",
+        raw="test output",
+        agent=agent.role
+    )):
+        # Run crew
+        crew.kickoff()
+    
+    # Verify agentops.end_session was called correctly
+    mock_agentops.end_session.assert_called_once_with(
+        end_state="Success",
+        end_state_reason="Finished Execution",
+        is_auto_end=True
+    )
+
 ceo = Agent(
     role="CEO",
     goal="Make sure the writers in your company produce amazing content.",