Fix autogen reload (#468)

* make start_session non blocking * fix build issue * bump version * callback for start session now that it is async * fix callback * unpin dependencies (#434) Co-authored-by: Howard Gil <[email protected]> * bump version number * wip * change autogen getting run every time * remove prints * remove more prints * suppress warnings * exponential retry to close data * removed event counter * fix requests mock * remove print * fixed more tests * removed inits from test_agent; does not require a client * create requests fixture * Scope fixtures * black run * remove bad files * add spaces back * revert session threading changes * remove callback * revert session changes * revert session * fix test * update tests * update tox to install local build instead of pypi * replace http client with requests so requests_mock works properly * fixed multiple sessions * fix tool recorder * removed test logs, fixed request count tests * set fixture scopes * Fixed missing async tests failing in tox, updated tox * fixed missing pass in tests * fixed timing * created rc branch --------- Co-authored-by: Shawn Qiu <[email protected]> Co-authored-by: Howard Gil <[email protected]>
AgentOps-AI · Nov 1, 2024 · f4d8347 · f4d8347
1 parent 6b37e1a
commit f4d8347
Show file tree

Hide file tree

Showing 14 changed files with 138 additions and 116 deletions.
diff --git a/agentops/__init__.py b/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)

diff --git a/agentops/client.py b/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
@@ -305,7 +302,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
 

diff --git a/agentops/helpers.py b/agentops/helpers.py
@@ -3,7 +3,7 @@
 from datetime import datetime, timezone
 import inspect
 from typing import Union
-import http.client
+import requests
 import json
 from importlib.metadata import version, PackageNotFoundError
 
@@ -127,25 +127,25 @@ def get_agentops_version():
 
 
 def check_agentops_update():
-    # using http.client to avoid this call being caught by requests_mock on tests
-    conn = http.client.HTTPSConnection("pypi.org")
-    conn.request("GET", "/pypi/agentops/json")
-    response = conn.getresponse()
-    data = response.read().decode()
-    json_data = json.loads(data)
+    try:
+        response = requests.get("https://pypi.org/pypi/agentops/json")
 
-    if response.status == 200:
-        latest_version = json_data["info"]["version"]
+        if response.status_code == 200:
+            json_data = response.json()
+            latest_version = json_data["info"]["version"]
 
-        try:
-            current_version = version("agentops")
-        except PackageNotFoundError:
-            return None
-
-        if not latest_version == current_version:
-            logger.warning(
-                f" WARNING: agentops is out of date. Please update with the command: 'pip install --upgrade agentops'"
-            )
+            try:
+                current_version = version("agentops")
+            except PackageNotFoundError:
+                return None
+
+            if not latest_version == current_version:
+                logger.warning(
+                    f" WARNING: agentops is out of date. Please update with the command: 'pip install --upgrade agentops'"
+                )
+    except Exception as e:
+        logger.debug(f"Failed to check for updates: {e}")
+        return None
 
 
 # Function decorator that prints function name and its arguments to the console for debug purposes

diff --git a/agentops/host_env.py b/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()
             }
         }

diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "agentops"
-version = "0.3.14"
+version = "0.3.15rc1"
 authors = [
   { name="Alex Reibman", email="[email protected]" },
   { name="Shawn Qiu", email="[email protected]" },
@@ -45,3 +45,7 @@ max_line_length = 120
 
 [project.scripts]
 agentops = "agentops.cli:main"
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"
diff --git a/tests/test_agent.py b/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"

diff --git a/tests/test_canary.py b/tests/test_canary.py
@@ -13,24 +13,25 @@ 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"
-        m.post(url + "/v2/create_events", text="ok")
+        m.post(url + "/v2/create_events", json={"status": "ok"})
         m.post(
             url + "/v2/create_session", json={"status": "success", "jwt": "some_jwt"}
         )
         m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", text="ok")
+        m.post(url + "/v2/developer_errors", json={"status": "ok"})
+        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
         yield m
 
 
 class TestCanary:
     def setup_method(self):
         self.url = "https://api.agentops.ai"
         self.api_key = "11111111-1111-4111-8111-111111111111"
-        agentops.init(api_key=self.api_key, max_wait_time=5, auto_start_session=False)
+        agentops.init(api_key=self.api_key, max_wait_time=500, auto_start_session=False)
 
     def test_agent_ops_record(self, mock_req):
         # Arrange
@@ -39,10 +40,11 @@ def test_agent_ops_record(self, mock_req):
 
         # Act
         agentops.record(ActionEvent(event_type))
-        time.sleep(0.1)
+        time.sleep(2)
+
+        # 3 requests: check_for_updates, create_session, create_events
+        assert len(mock_req.request_history) == 3
 
-        # Assert
-        assert len(mock_req.request_history) == 2
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
         assert request_json["events"][0]["event_type"] == event_type

diff --git a/tests/test_events.py b/tests/test_events.py
@@ -13,16 +13,21 @@ 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"
-        m.post(url + "/v2/create_events", text="ok")
+        m.post(url + "/v2/create_events", json={"status": "ok"})
         m.post(
             url + "/v2/create_session", json={"status": "success", "jwt": "some_jwt"}
         )
         m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", text="ok")
+        m.post(url + "/v2/developer_errors", json={"status": "ok"})
+        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
+
+        m.post(
+            url + "/v2/reauthorize_jwt", json={"status": "success", "jwt": "some_jwt"}
+        )
         yield m
 
 

diff --git a/tests/test_pre_init.py b/tests/test_pre_init.py
@@ -18,15 +18,19 @@ 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"
-        m.post(url + "/v2/create_agent", text="ok")
-        m.post(url + "/v2/update_session", text="ok")
+        m.post(url + "/v2/create_agent", json={"status": "success"})
+        m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
         m.post(
             url + "/v2/create_session", json={"status": "success", "jwt": "some_jwt"}
         )
+        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
+
+        m.post(url + "/v2/create_events", json={"status": "ok"})
+        m.post(url + "/v2/developer_errors", json={"status": "ok"})
 
         yield m
 
@@ -48,10 +52,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)
+
+        # 4 requests: check_for_updates, create_session, create_agent, update_session
+        assert len(mock_req.request_history) == 4
+
+        assert (
+            mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
+        )
+
+        mock_req.reset()
diff --git a/tests/test_record_action.py b/tests/test_record_action.py
@@ -18,11 +18,11 @@ 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"
-        m.post(url + "/v2/create_events", text="ok")
+        m.post(url + "/v2/create_events", json={"status": "ok"})
 
         # Use iter to create an iterator that can return the jwt values
         jwt_tokens = iter(jwts)
@@ -33,8 +33,10 @@ def create_session_response(request, context):
             return {"status": "success", "jwt": next(jwt_tokens)}
 
         m.post(url + "/v2/create_session", json=create_session_response)
+        m.post(url + "/v2/create_events", json={"status": "ok"})
         m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", text="ok")
+        m.post(url + "/v2/developer_errors", json={"status": "ok"})
+        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
 
         yield m
 
@@ -44,7 +46,7 @@ def setup_method(self):
         self.url = "https://api.agentops.ai"
         self.api_key = "11111111-1111-4111-8111-111111111111"
         self.event_type = "test_event_type"
-        agentops.init(self.api_key, max_wait_time=5, auto_start_session=False)
+        agentops.init(self.api_key, max_wait_time=50, auto_start_session=False)
 
     def test_record_action_decorator(self, mock_req):
         agentops.start_session()
@@ -57,8 +59,8 @@ def add_two(x, y):
         add_two(3, 4)
         time.sleep(0.1)
 
-        # Assert
-        assert len(mock_req.request_history) == 2
+        # 3 requests: check_for_updates, start_session, record_action
+        assert len(mock_req.request_history) == 3
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json["events"][0]["action_type"] == self.event_type
@@ -78,8 +80,8 @@ def add_two(x, y):
         add_two(3, 4)
         time.sleep(0.1)
 
-        # Assert
-        assert len(mock_req.request_history) == 2
+        # 3 requests: check_for_updates, start_session, record_action
+        assert len(mock_req.request_history) == 3
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json["events"][0]["action_type"] == "add_two"
@@ -98,14 +100,19 @@ def add_three(x, y, z=3):
 
         # Act
         add_three(1, 2)
-        time.sleep(0.1)
-        add_three(1, 2)
-        time.sleep(0.1)
+        add_three(1, 2, 4)
 
-        # Assert
+        time.sleep(1.5)
+
+        # 3 requests: check_for_updates, start_session, record_action
         assert len(mock_req.request_history) == 3
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
         request_json = mock_req.last_request.json()
+
+        assert request_json["events"][1]["action_type"] == self.event_type
+        assert request_json["events"][1]["params"] == {"x": 1, "y": 2, "z": 4}
+        assert request_json["events"][1]["returns"] == 7
+
         assert request_json["events"][0]["action_type"] == self.event_type
         assert request_json["events"][0]["params"] == {"x": 1, "y": 2, "z": 3}
         assert request_json["events"][0]["returns"] == 6
@@ -128,7 +135,7 @@ async def async_add(x, y):
         # Assert
         assert result == 7
         # Assert
-        assert len(mock_req.request_history) == 2
+        assert len(mock_req.request_history) == 3
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
         request_json = mock_req.last_request.json()
         assert request_json["events"][0]["action_type"] == self.event_type
@@ -156,11 +163,10 @@ def add_three(x, y, z=3):
         # Act
         add_three(1, 2, session=session_1)
         time.sleep(0.1)
-        add_three(1, 2, session=session_2)
+        add_three(1, 2, 3, session=session_2)
         time.sleep(0.1)
 
-        # Assert
-        assert len(mock_req.request_history) == 4
+        assert len(mock_req.request_history) == 5
 
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
@@ -207,7 +213,7 @@ async def async_add(x, y):
         time.sleep(0.1)
 
         # Assert
-        assert len(mock_req.request_history) == 4
+        assert len(mock_req.request_history) == 5
 
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key