handle calling create_agent before init

agentops/client.py

@@ -38,6 +38,7 @@ def __init__(self):
         self._llm_tracker: Optional[LlmTracker] = None
         self._sessions: List[Session] = active_sessions
         self._config = Configuration()
+        self._pre_init_queue = {"agents": []}
 
         self.configure(
             api_key=os.environ.get("AGENTOPS_API_KEY"),
@@ -106,6 +107,13 @@ def initialize(self) -> Union[Session, None]:
         if self._config.auto_start_session:
             session = self.start_session()
 
+        if session:
+            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"] = []
+
         return session
 
     def _initialize_partner_framework(self) -> None:
@@ -234,6 +242,13 @@ def start_session(
             config=self._config,
         )
 
+        if self._pre_init_queue["agents"] and 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"] = []
+
         if not session.is_running:
             return logger.error("Failed to start session")
 
@@ -294,7 +309,9 @@ def create_agent(
             # if no session passed, assume single session
             session = self._safe_get_session()
             if session is None:
-                return
+                self._pre_init_queue["agents"].append(
+                    {"name": name, "agent_id": agent_id}
+                )
             session.create_agent(name=name, agent_id=agent_id)
 
         return agent_id

agentops/decorators.py

@@ -326,11 +326,11 @@ def new_init(self, *args, **kwargs):
 
                     original_init(self, *args, **kwargs)
 
-                    if not Client().is_initialized:
-                        Client().add_pre_init_warning(
-                            f"Failed to track an agent {name} because agentops.init() was not "
-                            + "called before initializing the agent with the @track_agent decorator."
-                        )
+                    # if not Client().is_initialized:
+                    #     Client().add_pre_init_warning(
+                    #         f"Failed to track an agent {name} because agentops.init() was not "
+                    #         + "called before initializing the agent with the @track_agent decorator."
+                    #     )
 
                     self.agent_ops_agent_id = str(uuid4())
 
@@ -344,14 +344,15 @@ def new_init(self, *args, **kwargs):
                         session=session,
                     )
                 except AttributeError as e:
-                    Client().add_pre_init_warning(
-                        f"Failed to track an agent {name} because agentops.init() was not "
-                        + "called before initializing the agent with the @track_agent decorator."
-                    )
-                    logger.warning(
-                        "Failed to track an agent. This often happens if agentops.init() was not "
-                        "called before initializing an agent with the @track_agent decorator."
-                    )
+                    print(e)
+                    # Client().add_pre_init_warning(
+                    #     f"Failed to track an agent {name} because agentops.init() was not "
+                    #     + "called before initializing the agent with the @track_agent decorator."
+                    # )
+                    # logger.warning(
+                    #     "Failed to track an agent. This often happens if agentops.init() was not "
+                    #     "called before initializing an agent with the @track_agent decorator."
+                    # )
                     original_init(self, *args, **kwargs)
 
             obj.__init__ = new_init

tests/test_pre_init.py

@@ -0,0 +1,57 @@
+import pytest
+import requests_mock
+import time
+import agentops
+from agentops import record_action, track_agent
+from datetime import datetime
+from agentops.singleton import clear_singletons
+import contextlib
+
+jwts = ["some_jwt", "some_jwt2", "some_jwt3"]
+
+
+@pytest.fixture(autouse=True)
+def setup_teardown():
+    clear_singletons()
+    yield
+    agentops.end_all_sessions()  # teardown part
+
+
+@contextlib.contextmanager
+@pytest.fixture(autouse=True)
+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_session", json={"status": "success", "jwt": "some_jwt"}
+        )
+
+        yield m
+
+
+@track_agent(name="TestAgent")
+class BasicAgent:
+    def __init__(self):
+        pass
+
+
+class TestPreInit:
+    def setup_method(self):
+        self.url = "https://api.agentops.ai"
+        self.api_key = "11111111-1111-4111-8111-111111111111"
+
+    def test_track_agent(self, mock_req):
+        agent = BasicAgent()
+
+        assert len(mock_req.request_history) == 0
+
+        agentops.init(api_key=self.api_key)
+
+        # 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")