diff --git a/agentops/__init__.py b/agentops/__init__.py index 4c8f8dae..ccb5427d 100755 --- 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 index 0c41f0c9..0942f1e2 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -87,7 +87,6 @@ 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." @@ -95,8 +94,6 @@ def initialize(self) -> Union[Session, None]: ) 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 diff --git a/agentops/host_env.py b/agentops/host_env.py index d446d347..e466c86e 100644 --- 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/agentops/http_client.py b/agentops/http_client.py index 32d7d336..4d18274b 100644 --- a/agentops/http_client.py +++ b/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 diff --git a/agentops/session.py b/agentops/session.py index 6fd51e61..ab3a833f 100644 --- a/agentops/session.py +++ b/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): diff --git a/tests/test_agent.py b/tests/test_agent.py index 408161b9..9005dc75 100644 --- 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 index 8ff49a01..4d430ed7 100644 --- a/tests/test_canary.py +++ b/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" diff --git a/tests/test_events.py b/tests/test_events.py index a2d70d75..bf1c22d7 100644 --- a/tests/test_events.py +++ b/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" diff --git a/tests/test_pre_init.py b/tests/test_pre_init.py index f87219ac..be4b7bdd 100644 --- a/tests/test_pre_init.py +++ b/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() diff --git a/tests/test_record_action.py b/tests/test_record_action.py index 058e13dd..2ca69aa5 100644 --- a/tests/test_record_action.py +++ b/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" diff --git a/tests/test_record_tool.py b/tests/test_record_tool.py index 29a5d473..f47f07a7 100644 --- a/tests/test_record_tool.py +++ b/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" diff --git a/tests/test_session.py b/tests/test_session.py index 075aca84..46be74cd 100644 --- a/tests/test_session.py +++ b/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" diff --git a/tests/test_teardown.py b/tests/test_teardown.py index b86ea079..f9dbe22b 100644 --- a/tests/test_teardown.py +++ b/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"