diff --git a/agentops/client.py b/agentops/client.py index cbb783371..48255a750 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -113,7 +113,7 @@ def record(self, event: Event | ErrorEvent): self._worker.add_event(event.__dict__) else: logging.warning( - "AgentOps: Cannot record event - no current session") + "🖇 AgentOps: Cannot record event - no current session") def _record_event_sync(self, func, event_name, *args, **kwargs): init_time = get_ISO_time() @@ -205,17 +205,17 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi config: (Configuration, optional): Client configuration object """ if self._session is not None: - return logging.warning("AgentOps: Cannot start session - session already started") + return logging.warning("🖇 AgentOps: Cannot start session - session already started") if not config and not self.config: - return logging.warning("AgentOps: Cannot start session - missing configuration") + return logging.warning("🖇 AgentOps: Cannot start session - missing configuration") self._session = Session(uuid4(), tags or self._tags, host_env=get_host_env()) self._worker = Worker(config or self.config) start_session_result = self._worker.start_session(self._session) if not start_session_result: self._session = None - return logging.warning("AgentOps: Cannot start session") + return logging.warning("🖇 AgentOps: Cannot start session") logging.info('View info on this session at https://app.agentops.ai/drilldown?session_id={}' .format(self._session.session_id)) @@ -233,14 +233,15 @@ def end_session(self, video (str, optional): The video screen recording of the session """ if self._session is None or self._session.has_ended: - return logging.warning("AgentOps: Cannot end session - no current session") + return logging.warning("🖇 AgentOps: Cannot end session - no current session") if not any(end_state == state.value for state in EndState): - return logging.warning("AgentOps: Invalid end_state. Please use one of the EndState enums") + return logging.warning("🖇 AgentOps: Invalid end_state. Please use one of the EndState enums") self._session.video = video self._session.end_session(end_state, end_state_reason) - self._worker.end_session(self._session) + token_cost = float(self._worker.end_session(self._session)) + print('🖇 AgentOps: This run cost ${:.6f}'.format(token_cost)) self._session = None self._worker = None @@ -265,7 +266,7 @@ def signal_handler(signum, frame): """ signal_name = 'SIGINT' if signum == signal.SIGINT else 'SIGTERM' logging.info( - f'AgentOps: {signal_name} detected. Ending session...') + f'🖇 AgentOps: {signal_name} detected. Ending session...') self.end_session(end_state='Fail', end_state_reason=f'Signal {signal_name} detected') sys.exit(0) diff --git a/agentops/config.py b/agentops/config.py index 2bd562627..f5751b24e 100644 --- a/agentops/config.py +++ b/agentops/config.py @@ -35,7 +35,7 @@ def __init__(self, if not api_key: api_key = environ.get('AGENTOPS_API_KEY', None) if not api_key: - raise ConfigurationError("AgentOps: No API key provided - no data will be recorded.") + raise ConfigurationError("🖇 AgentOps: No API key provided - no data will be recorded.") if not parent_key: parent_key = environ.get('AGENTOPS_PARENT_KEY', None) diff --git a/agentops/http_client.py b/agentops/http_client.py index bf404d9e3..f1d205f06 100644 --- a/agentops/http_client.py +++ b/agentops/http_client.py @@ -83,7 +83,7 @@ def post(url: str, payload: bytes, api_key: Optional[str] = None, parent_key: Op result.code = 408 result.status = HttpStatus.TIMEOUT logging.warning( - 'AgentOps: Could not post data - connection timed out') + '🖇 AgentOps: Could not post data - connection timed out') except requests.exceptions.HTTPError as e: try: result.parse(e.response) @@ -97,11 +97,11 @@ def post(url: str, payload: bytes, api_key: Optional[str] = None, parent_key: Op if result.code == 401: logging.warning( - f'AgentOps: Could not post data - API server rejected your API key: {api_key}') + f'🖇 AgentOps: Could not post data - API server rejected your API key: {api_key}') if result.code == 400: - logging.warning(f'AgentOps: Could not post data - {result.body}') + logging.warning(f'🖇 AgentOps: Could not post data - {result.body}') if result.code == 500: logging.warning( - f'AgentOps: Could not post data - internal server error') + f'🖇 AgentOps: Could not post data - internal server error') return result diff --git a/agentops/llm_tracker.py b/agentops/llm_tracker.py index e2ef73278..1aefc392e 100644 --- a/agentops/llm_tracker.py +++ b/agentops/llm_tracker.py @@ -61,7 +61,7 @@ def handle_stream_chunk(chunk): self.client.record(ErrorEvent(trigger_event=self.llm_event, details={f"{type(e).__name__}": str(e)})) # TODO: This error is specific to only one path of failure. Should be more generic or have different logging for different paths logging.warning( - f"AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") + f"🖇 AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") # if the response is a generator, decorate the generator if inspect.isasyncgen(response): @@ -100,7 +100,7 @@ def generator(): self.client.record(ErrorEvent(trigger_event=self.llm_event, details={f"{type(e).__name__}": str(e)})) # TODO: This error is specific to only one path of failure. Should be more generic or have different logging for different paths logging.warning( - f"AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") + f"🖇 AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") return response @@ -146,7 +146,7 @@ def handle_stream_chunk(chunk: ChatCompletionChunk): self.client.record(ErrorEvent(trigger_event=self.llm_event, details={f"{type(e).__name__}": str(e)})) # TODO: This error is specific to only one path of failure. Should be more generic or have different logging for different paths logging.warning( - f"AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") + f"🖇 AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") # if the response is a generator, decorate the generator if isinstance(response, Stream): @@ -191,7 +191,7 @@ async def async_generator(): self.client.record(ErrorEvent(trigger_event=self.llm_event, details={f"{type(e).__name__}": str(e)})) # TODO: This error is specific to only one path of failure. Should be more generic or have different logging for different paths logging.warning( - f"AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") + f"🖇 AgentOps: Unable to parse a chunk for LLM call {kwargs} - skipping upload to AgentOps") return response diff --git a/agentops/meta_client.py b/agentops/meta_client.py index 29f785886..d90d80356 100644 --- a/agentops/meta_client.py +++ b/agentops/meta_client.py @@ -45,7 +45,7 @@ def wrapper(self, *args, **kwargs): try: return method(self, *args, **kwargs) except Exception as e: - logging.warning(f"AgentOps: Error: {e}") + logging.warning(f"🖇 AgentOps: Error: {e}") config = getattr(self, 'config', None) if config is not None: type(self).send_exception_to_server(e, self.config._api_key) diff --git a/agentops/worker.py b/agentops/worker.py index 6b7320efe..8d9d3f93b 100644 --- a/agentops/worker.py +++ b/agentops/worker.py @@ -59,7 +59,7 @@ def start_session(self, session: Session) -> None: return True - def end_session(self, session: Session) -> None: + def end_session(self, session: Session) -> str: self.stop_flag.set() self.thread.join(timeout=1) self.flush_queue() @@ -70,12 +70,14 @@ def end_session(self, session: Session) -> None: "session": session.__dict__ } - HttpClient.post(f'{self.config.endpoint}/sessions', + res = HttpClient.post(f'{self.config.endpoint}/sessions', json.dumps(filter_unjsonable( payload)).encode("utf-8"), self.config.api_key, self.config.parent_key) + return res.body.get('token_cost', "unknown") + def update_session(self, session: Session) -> None: with self.lock: payload = { diff --git a/examples/openai-gpt.ipynb b/examples/openai-gpt.ipynb index 5ee44ad43..f2599d935 100644 --- a/examples/openai-gpt.ipynb +++ b/examples/openai-gpt.ipynb @@ -220,15 +220,15 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "4ca2b49fc06adddb", - "metadata": { - "collapsed": false - }, "outputs": [], "source": [ "agentops.end_session('Success')" - ] + ], + "metadata": { + "collapsed": false + }, + "id": "4ca2b49fc06adddb", + "execution_count": null } ], "metadata": { diff --git a/tests/test_canary.py b/tests/test_canary.py index e960c726c..09c4c30be 100644 --- a/tests/test_canary.py +++ b/tests/test_canary.py @@ -1,4 +1,5 @@ import pytest +import requests import requests_mock import time import agentops @@ -10,19 +11,15 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' m.post(url + '/events', text='ok') - m.post(url + '/sessions', text='ok') + m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) yield m - class TestCanary: def setup_method(self): self.url = 'https://api.agentops.ai' self.api_key = "random_api_key" agentops.init(api_key=self.api_key, max_wait_time=5, auto_start_session=False) - def teardown_method(self): - agentops.end_session(end_state='Success') - def test_agent_ops_record(self, mock_req): # Arrange event_type = 'test_event_type' @@ -37,3 +34,5 @@ def test_agent_ops_record(self, mock_req): request_json = mock_req.last_request.json() assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key assert request_json['events'][0]['event_type'] == event_type + + agentops.end_session('Success') diff --git a/tests/test_record_function.py b/tests/test_record_function.py index de03a81f3..5896381b3 100644 --- a/tests/test_record_function.py +++ b/tests/test_record_function.py @@ -10,7 +10,7 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' m.post(url + '/events', text='ok') - m.post(url + '/sessions', text='ok') + m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) yield m class TestRecordAction: @@ -19,12 +19,10 @@ def setup_method(self): self.api_key = "random_api_key" self.event_type = 'test_event_type' agentops.init(self.api_key, max_wait_time=5, auto_start_session=False) - agentops.start_session() - - def teardown_method(self): - agentops.end_session(end_state='Success') def test_record_function_decorator(self, mock_req): + agentops.start_session() + @record_function(event_name=self.event_type) def add_two(x, y): return x + y @@ -34,14 +32,18 @@ def add_two(x, y): time.sleep(0.1) # Assert - assert len(mock_req.request_history) == 1 + assert len(mock_req.request_history) == 2 assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key request_json = mock_req.last_request.json() assert request_json['events'][0]['action_type'] == self.event_type assert request_json['events'][0]['params'] == {'x': 3, 'y': 4} assert request_json['events'][0]['returns'] == 7 + agentops.end_session(end_state='Success') + def test_record_function_decorator_multiple(self, mock_req): + agentops.start_session() + # Arrange @record_function(event_name=self.event_type) def add_three(x, y, z=3): @@ -54,15 +56,18 @@ def add_three(x, y, z=3): time.sleep(0.1) # Assert - assert len(mock_req.request_history) == 2 + assert len(mock_req.request_history) == 3 assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key request_json = mock_req.last_request.json() 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 + agentops.end_session(end_state='Success') + @pytest.mark.asyncio async def test_async_function_call(self, mock_req): + agentops.start_session() @record_function(self.event_type) async def async_add(x, y): @@ -76,7 +81,7 @@ async def async_add(x, y): # Assert assert result == 7 # Assert - assert len(mock_req.request_history) == 1 + assert len(mock_req.request_history) == 2 assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key request_json = mock_req.last_request.json() assert request_json['events'][0]['action_type'] == self.event_type @@ -87,4 +92,6 @@ async def async_add(x, y): end = datetime.fromisoformat( request_json['events'][0]['end_timestamp'].replace('Z', '+00:00')) - assert (end - init).total_seconds() >= 0.1 \ No newline at end of file + assert (end - init).total_seconds() >= 0.1 + + agentops.end_session(end_state='Success') \ No newline at end of file diff --git a/tests/test_session.py b/tests/test_session.py index 2efe82a4b..339835de7 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -10,7 +10,7 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' m.post(url + '/events', text='ok') - m.post(url + '/sessions', text='ok') + m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) yield m diff --git a/tests/test_teardown.py b/tests/test_teardown.py index 06bcceb4b..114a6aff3 100644 --- a/tests/test_teardown.py +++ b/tests/test_teardown.py @@ -10,7 +10,7 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' m.post(url + '/events', text='ok') - m.post(url + '/sessions', text='ok') + m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) yield m