print token cost on session end (#145)

* print token cost on session end

* dupe paperclip

* responses

* tests work

---------

Co-authored-by: Shawn Qiu <siyangqiu@gmail.com>

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)

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)

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

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
 

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)

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 = {

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": {

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')

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
+        assert (end - init).total_seconds() >= 0.1
+
+        agentops.end_session(end_state='Success')

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
 
 

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