Copy of PR #435 - make start_session non blocking

agentops/__init__.py

@@ -22,6 +22,7 @@
 
 if "autogen" in sys.modules:
     Client().configure(instrument_llm_calls=False)
+    Client()._initialize_autogen_logger()
     Client().add_default_tags(["autogen"])
 
 if "crewai" in sys.modules:
@@ -72,7 +73,6 @@ def init(
     Client().unsuppress_logs()
     t = threading.Thread(target=check_agentops_update)
     t.start()
-
     if Client().is_initialized:
         return logger.warning(
             "AgentOps has already been initialized. If you are trying to start a session, call agentops.start_session() instead."
@@ -101,7 +101,6 @@ def init(
                 "auto_start_session is set to False - inherited_session_id will not be used to automatically start a session"
             )
             return Client().initialize()
-
         Client().configure(auto_start_session=False)
         Client().initialize()
         return Client().start_session(inherited_session_id=inherited_session_id)

agentops/client.py

@@ -87,16 +87,13 @@ def initialize(self) -> Union[Session, None]:
             return
 
         self.unsuppress_logs()
-
         if self._config.api_key is None:
             return logger.error(
                 "Could not initialize AgentOps client - API Key is missing."
                 + "\n\t    Find your API key at https://app.agentops.ai/settings/projects"
             )
 
         self._handle_unclean_exits()
-        self._initialize_partner_framework()
-
         self._initialized = True
 
         if self._config.instrument_llm_calls:
@@ -116,7 +113,7 @@ def initialize(self) -> Union[Session, None]:
 
         return session
 
-    def _initialize_partner_framework(self) -> None:
+    def _initialize_autogen_logger(self) -> None:
         try:
             import autogen
             from .partners.autogen_logger import AutogenLogger
@@ -235,11 +232,30 @@ def start_session(
         if tags is not None:
             session_tags.update(tags)
 
+        def _start_session_callback(session: Session):
+            if session.is_running:
+                if len(self._pre_init_queue["agents"]) > 0:
+                    for agent_args in self._pre_init_queue["agents"]:
+                        session.create_agent(
+                            name=agent_args["name"], agent_id=agent_args["agent_id"]
+                        )
+                    self._pre_init_queue["agents"] = []
+
+                logger.info(
+                    colored(
+                        f"\x1b[34mSession Replay: https://app.agentops.ai/drilldown?session_id={session.session_id}\x1b[0m",
+                        "blue",
+                    )
+                )
+            else:
+                self._sessions.remove(session)
+
         session = Session(
             session_id=session_id,
             tags=list(session_tags),
             host_env=get_host_env(self._config.env_data_opt_out),
             config=self._config,
+            callback=_start_session_callback,
         )
 
         if not session.is_running:
@@ -305,7 +321,8 @@ def create_agent(
                 self._pre_init_queue["agents"].append(
                     {"name": name, "agent_id": agent_id}
                 )
-            session.create_agent(name=name, agent_id=agent_id)
+            else:
+                session.create_agent(name=name, agent_id=agent_id)
 
         return agent_id
 

agentops/host_env.py

@@ -47,13 +47,11 @@ def get_sys_packages():
 
 
 def get_installed_packages():
-
     try:
         return {
-            # TODO: test
             # TODO: add to opt out
             "Installed Packages": {
-                dist.metadata["Name"]: dist.version
+                dist.metadata.get("Name"): dist.metadata.get("Version")
                 for dist in importlib.metadata.distributions()
             }
         }

agentops/http_client.py

@@ -81,11 +81,9 @@ def post(
 
             if jwt is not None:
                 JSON_HEADER["Authorization"] = f"Bearer {jwt}"
-
             res = request_session.post(
                 url, data=payload, headers=JSON_HEADER, timeout=20
             )
-
             result.parse(res)
         except requests.exceptions.Timeout:
             result.code = 408

agentops/session.py

@@ -5,7 +5,7 @@
 import time
 from decimal import ROUND_HALF_UP, Decimal
 from termcolor import colored
-from typing import Optional, List, Union
+from typing import Optional, List, Union, Callable
 from uuid import UUID, uuid4
 from datetime import datetime
 
@@ -40,6 +40,7 @@ def __init__(
         config: Configuration,
         tags: Optional[List[str]] = None,
         host_env: Optional[dict] = None,
+        callback: Optional[Callable[["Session"], None]] = None,
     ):
         self.end_timestamp = None
         self.end_state: Optional[str] = None
@@ -60,17 +61,13 @@ def __init__(
             "errors": 0,
             "apis": 0,
         }
-
+        self.is_running = False
+        active_sessions.append(self)
         self.stop_flag = threading.Event()
-        self.thread = threading.Thread(target=self._run)
+        self.thread = threading.Thread(target=self._run, args=(callback,))
         self.thread.daemon = True
         self.thread.start()
 
-        self.is_running = self._start_session()
-        if self.is_running == False:
-            self.stop_flag.set()
-            self.thread.join(timeout=1)
-
     def set_video(self, video: str) -> None:
         """
         Sets a url to the video recording of the session.
@@ -274,16 +271,21 @@ def _start_session(self):
                     self.config.parent_key,
                 )
             except ApiServerException as e:
-                return logger.error(f"Could not start session - {e}")
+                logger.error(f"Could not start session - {e}")
+                return False
 
             logger.debug(res.body)
 
             if res.code != 200:
+                logger.error(f"Could not start session - server error")
                 return False
 
             jwt = res.body.get("jwt", None)
             self.jwt = jwt
             if jwt is None:
+                logger.error(
+                    f"Could not start session - server could not authenticate your API Key"
+                )
                 return False
 
             session_url = res.body.get(
@@ -359,10 +361,19 @@ def _flush_queue(self) -> None:
                     elif event_type == "apis":
                         self.event_counts["apis"] += 1
 
-    def _run(self) -> None:
+    def _run(self, callback: Optional[Callable[["Session"], None]] = None) -> None:
+        self.is_running = self._start_session()
+
+        if callback:
+            callback(self)
+
+        if self.is_running == False:
+            self.stop_flag.set()
+            self.thread.join(timeout=1)
+
         while not self.stop_flag.is_set():
             time.sleep(self.config.max_wait_time / 1000)
-            if self.queue:
+            if self.queue and self.jwt is not None:
                 self._flush_queue()
 
     def create_agent(self, name, agent_id):

tests/test_agent.py

@@ -1,13 +1,9 @@
 from unittest import TestCase
-
 from agentops import track_agent
-import agentops
 
 
 class TrackAgentTests(TestCase):
     def test_track_agent_with_class(self):
-        agentops.init()
-
         @track_agent(name="agent_name")
         class TestAgentClass:
             t = "a"
@@ -19,8 +15,6 @@ class TestAgentClass:
         self.assertIsNotNone(getattr(obj, "agent_ops_agent_id"))
 
     def test_track_agent_with_class_name(self):
-        agentops.init()
-
         @track_agent(name="agent_name")
         class TestAgentClass:
             t = "a"

tests/test_canary.py

@@ -13,7 +13,7 @@ def setup_teardown():
     agentops.end_all_sessions()  # teardown part
 
 
-@pytest.fixture
+@pytest.fixture(autouse=True, scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"

tests/test_events.py

@@ -13,7 +13,7 @@ def setup_teardown():
     agentops.end_all_sessions()  # teardown part
 
 
-@pytest.fixture
+@pytest.fixture(scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"

tests/test_pre_init.py

@@ -18,7 +18,7 @@ def setup_teardown():
 
 
 @contextlib.contextmanager
-@pytest.fixture(autouse=True)
+@pytest.fixture(autouse=True, scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"
@@ -48,10 +48,20 @@ def test_track_agent(self, mock_req):
         assert len(mock_req.request_history) == 0
 
         agentops.init(api_key=self.api_key)
+        time.sleep(1)
 
         # Assert
         # start session and create agent
-        assert len(mock_req.request_history) == 2
-        assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-
         agentops.end_session(end_state="Success")
+
+        # Wait for flush
+        time.sleep(1.5)
+
+        # 3 requests: create session, create agent, update session
+        assert len(mock_req.request_history) == 3
+
+        assert (
+            mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
+        )
+
+        mock_req.reset()

tests/test_record_action.py

@@ -18,7 +18,7 @@ def setup_teardown():
 
 
 @contextlib.contextmanager
-@pytest.fixture(autouse=True)
+@pytest.fixture(autouse=True, scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"

tests/test_record_tool.py

@@ -19,7 +19,7 @@ def setup_teardown():
 
 
 @contextlib.contextmanager
-@pytest.fixture(autouse=True)
+@pytest.fixture(autouse=True, scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"

tests/test_session.py

@@ -13,7 +13,7 @@ def setup_teardown():
     agentops.end_all_sessions()  # teardown part
 
 
-@pytest.fixture
+@pytest.fixture(scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"

tests/test_teardown.py

@@ -3,7 +3,7 @@
 from agentops import Client
 
 
-@pytest.fixture
+@pytest.fixture(scope="function")
 def mock_req():
     with requests_mock.Mocker() as m:
         url = "https://api.agentops.ai"