Return token cost on session_end (#264)

* Return token cost on session_end

* Add type hint for return

* Add docstring

* Fix black formatting error

agentops/__init__.py

@@ -118,7 +118,7 @@ def end_session(
         video (str, optional): URL to a video recording of the session
         is_auto_end (bool, optional): is this an automatic use of end_session and should be skipped with bypass_auto_end_session
     """
-    Client().end_session(
+    return Client().end_session(
         end_state=end_state,
         end_state_reason=end_state_reason,
         video=video,

agentops/client.py

@@ -1,8 +1,8 @@
 """
-    AgentOps client module that provides a client class with public interfaces and configuration.
+AgentOps client module that provides a client class with public interfaces and configuration.
 
-    Classes:
-        Client: Provides methods to interact with the AgentOps service.
+Classes:
+    Client: Provides methods to interact with the AgentOps service.
 """
 
 import os
@@ -80,7 +80,6 @@ def __init__(
         inherited_session_id: Optional[str] = None,
         skip_auto_end_session: Optional[bool] = False,
     ):
-
         if override is not None:
             logger.warning(
                 "The 'override' parameter is deprecated. Use 'instrument_llm_calls' instead.",
@@ -389,7 +388,7 @@ def end_session(
         end_state_reason: Optional[str] = None,
         video: Optional[str] = None,
         is_auto_end: Optional[bool] = None,
-    ):
+    ) -> Decimal:
         """
         End the current session with the AgentOps service.
 
@@ -398,6 +397,9 @@ def end_session(
             end_state_reason (str, optional): The reason for ending the session.
             video (str, optional): The video screen recording of the session
             is_auto_end (bool, optional): is this an automatic use of end_session and should be skipped with skip_auto_end_session
+
+        Returns:
+            Decimal: The token cost of the session. Returns 0 if the cost is unknown.
         """
 
         if is_auto_end and self.config.skip_auto_end_session:
@@ -420,6 +422,7 @@ def end_session(
 
         if token_cost is None or token_cost == "unknown":
             logger.info("Could not determine cost of run.")
+            token_cost_d = Decimal(0)
         else:
             token_cost_d = Decimal(token_cost)
             logger.info(
@@ -439,6 +442,7 @@ def end_session(
 
         self._session = None
         self._worker = None
+        return token_cost_d
 
     def create_agent(self, name: str, agent_id: Optional[str] = None):
         if agent_id is None:

examples/openai-gpt.ipynb

@@ -127,7 +127,7 @@
     "for chunk in res:\n",
     "    if chunk.choices[0].delta.content is not None:\n",
     "        full_response += chunk.choices[0].delta.content\n",
-    "    \n",
+    "\n",
     "print(full_response.strip())"
    ]
   },