From 69f856715f49e0096dd2e19b92a44dcd578e0843 Mon Sep 17 00:00:00 2001
From: Siddharth Prajosh <sprajosh@gmail.com>
Date: Thu, 27 Jun 2024 03:02:46 +0530
Subject: [PATCH] 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      |  2 +-
 agentops/client.py        | 14 +++++++++-----
 examples/openai-gpt.ipynb |  2 +-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/agentops/__init__.py b/agentops/__init__.py
index e8600cce..7a31698d 100755
--- a/agentops/__init__.py
+++ b/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,
diff --git a/agentops/client.py b/agentops/client.py
index 7108e393..34db1359 100644
--- a/agentops/client.py
+++ b/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:
diff --git a/examples/openai-gpt.ipynb b/examples/openai-gpt.ipynb
index ddc511c1..0fe4f241 100644
--- a/examples/openai-gpt.ipynb
+++ b/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())"
    ]
   },