From bb3b056e3e439b1e04a2827a16f765978d52548d Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 18 Dec 2024 04:45:05 +0000
Subject: [PATCH] refactor: Update end_session to return detailed statistics

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>
---
 agentops/client.py  | 14 ++++++--------
 agentops/session.py | 10 ++++++----
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/agentops/client.py b/agentops/client.py
index fb3e1793..a4814fb6 100644
--- a/agentops/client.py
+++ b/agentops/client.py
@@ -15,7 +15,7 @@
 import traceback
 from decimal import Decimal
 from functools import cached_property
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 from uuid import UUID, uuid4
 
 from termcolor import colored
@@ -247,7 +247,7 @@ def end_session(
         end_state_reason: Optional[str] = None,
         video: Optional[str] = None,
         is_auto_end: Optional[bool] = None,
-    ) -> Optional[Decimal]:
+    ) -> Optional[Dict[str, Any]]:
         """
         End the current session with the AgentOps service.
 
@@ -258,17 +258,15 @@ def end_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.
+            Dict[str, Any]: Session statistics including duration, cost, and event counts.
         """
         session = self._safe_get_session()
         if session is None:
-            return
+            return None
         if is_auto_end and self._config.skip_auto_end_session:
-            return
-
-        token_cost = session.end_session(end_state=end_state, end_state_reason=end_state_reason, video=video)
+            return None
 
-        return token_cost
+        return session.end_session(end_state=end_state, end_state_reason=end_state_reason, video=video)
 
     def create_agent(
         self,
diff --git a/agentops/session.py b/agentops/session.py
index 535bc997..69c55ce8 100644
--- a/agentops/session.py
+++ b/agentops/session.py
@@ -284,7 +284,7 @@ def end_session(
         end_state: str = "Indeterminate",
         end_state_reason: Optional[str] = None,
         video: Optional[str] = None,
-    ) -> Union[Decimal, None]:
+    ) -> Union[Dict[str, Any], None]:
         with self._end_session_lock:
             if not self.is_running:
                 return None
@@ -324,8 +324,9 @@ def end_session(
                     finally:
                         del self._span_processor
 
-                # 5. Final session update
-                if not (analytics_stats := self.get_analytics()):
+                # 5. Final session update and get analytics
+                analytics_stats = self.get_analytics()
+                if not analytics_stats:
                     return None
 
                 analytics = (
@@ -341,6 +342,7 @@ def end_session(
 
             except Exception as e:
                 logger.exception(f"Error during session end: {e}")
+                return None
             finally:
                 active_sessions.remove(self)  # First thing, get rid of the session
 
@@ -350,7 +352,7 @@ def end_session(
                         "blue",
                     )
                 )
-            return self.token_cost
+            return analytics_stats
 
     def add_tags(self, tags: List[str]) -> None:
         """