diff --git a/agentops/__init__.py b/agentops/__init__.py index 2edf10f8..1335f6a9 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -1,4 +1,5 @@ # agentops/__init__.py +import functools import os import logging from typing import Optional, List, Union @@ -18,6 +19,23 @@ except ModuleNotFoundError: pass +is_initialized = False + + +def noop(*args, **kwargs): + return + + +def check_init(child_function): + @functools.wraps(child_function) + def wrapper(*args, **kwargs): + if is_initialized: + return child_function(*args, **kwargs) + else: + return noop(*args, **kwargs) + + return wrapper + def init( api_key: Optional[str] = None, @@ -79,6 +97,9 @@ def init( skip_auto_end_session=skip_auto_end_session, ) + global is_initialized + is_initialized = True + return inherited_session_id or c.current_session_id @@ -118,9 +139,19 @@ def start_session( e.g. ["test_run"]. config: (Configuration, optional): Client configuration object """ - return Client().start_session(tags, config, inherited_session_id) + + try: + sess_result = Client().start_session(tags, config, inherited_session_id) + + global is_initialized + is_initialized = True + + return sess_result + except Exception: + pass +@check_init def record(event: Union[Event, ErrorEvent]): """ Record an event with the AgentOps service. @@ -131,6 +162,7 @@ def record(event: Union[Event, ErrorEvent]): Client().record(event) +@check_init def add_tags(tags: List[str]): """ Append to session tags at runtime. @@ -141,6 +173,7 @@ def add_tags(tags: List[str]): Client().add_tags(tags) +@check_init def set_tags(tags: List[str]): """ Replace session tags at runtime. @@ -169,5 +202,6 @@ def stop_instrumenting(): Client().stop_instrumenting() +@check_init def create_agent(name: str, agent_id: Optional[str] = None): return Client().create_agent(name=name, agent_id=agent_id) diff --git a/agentops/client.py b/agentops/client.py index 57f9cb60..7108e393 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -418,7 +418,7 @@ def end_session( self._session.end_session(end_state, end_state_reason) token_cost = self._worker.end_session(self._session) - if token_cost == "unknown": + if token_cost is None or token_cost == "unknown": logger.info("Could not determine cost of run.") else: token_cost_d = Decimal(token_cost) diff --git a/pyproject.toml b/pyproject.toml index 8ea93afc..4324f5c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "agentops" -version = "0.2.2" +version = "0.2.3" authors = [ { name="Alex Reibman", email="areibman@gmail.com" }, { name="Shawn Qiu", email="siyangqiu@gmail.com" }, diff --git a/tests/test_session.py b/tests/test_session.py index e8a1bf14..7148a008 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -29,7 +29,6 @@ def setup_method(self): def test_session(self, mock_req): agentops.start_session(config=self.config) - print(self.config.api_key) agentops.record(ActionEvent(self.event_type)) agentops.record(ActionEvent(self.event_type))